From 2f32af57ef0ec22a91e4c446d7b9ffea0524a090 Mon Sep 17 00:00:00 2001
From: MohammadReza Hosseini <mrhosseini1367@gmail.com>
Date: Sun, 21 Jun 2020 17:45:51 +0430
Subject: [PATCH] added support for both reorg and reorg3d layers of
 AlexeyAB/Darknet

---
 include/tkDNN/Layer.h             |  3 ++-
 include/tkDNN/pluginsRT/ReorgRT.h | 16 ++++++++++++----
 src/DarknetParser.cpp             |  2 +-
 src/NetworkRT.cpp                 |  2 +-
 src/Reorg.cpp                     |  9 ++++++---
 5 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h
index 1380c1d..acec247 100644
--- a/include/tkDNN/Layer.h
+++ b/include/tkDNN/Layer.h
@@ -534,13 +534,14 @@ public:
 class Reorg : public Layer {
 
 public:
-    Reorg(Network *net, int stride);
+    Reorg(Network *net, int stride, bool reorg3d = false);
     virtual ~Reorg();
     virtual layerType_t getLayerType() { return LAYER_REORG; };
 
     virtual dnnType* infer(dataDim_t &dim, dnnType* srcData);
 
     int stride;
+    bool reorg3d;
 };
 
 /**
diff --git a/include/tkDNN/pluginsRT/ReorgRT.h b/include/tkDNN/pluginsRT/ReorgRT.h
index afd205d..9a125f0 100644
--- a/include/tkDNN/pluginsRT/ReorgRT.h
+++ b/include/tkDNN/pluginsRT/ReorgRT.h
@@ -4,8 +4,9 @@
 class ReorgRT : public IPlugin {
 
 public:
-	ReorgRT(int stride) {
+	ReorgRT(int stride, bool reorg3d = false) {
 		this->stride = stride;
+		this->reorg3d = reorg3d;
 	}
 
 	~ReorgRT(){
@@ -21,9 +22,15 @@ public:
 	}
 
 	void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override {
-		c = outputDims[0].d[0];
-		h = outputDims[0].d[1];
-		w = outputDims[0].d[2];
+		if (reorg3d) {
+			c = outputDims[0].d[0];
+			h = outputDims[0].d[1];
+			w = outputDims[0].d[2];
+		} else {
+			c = inputDims[0].d[0];
+			h = inputDims[0].d[1];
+			w = inputDims[0].d[2];
+		}
 	}
 
 	int initialize() override {
@@ -60,4 +67,5 @@ public:
 	}
 
 	int c, h, w, stride;
+	bool reorg3d;
 };
diff --git a/src/DarknetParser.cpp b/src/DarknetParser.cpp
index 1626857..e101136 100644
--- a/src/DarknetParser.cpp
+++ b/src/DarknetParser.cpp
@@ -153,7 +153,7 @@ namespace tk { namespace dnn {
             netLayers.push_back(new tk::dnn::Route(net, layers.data(), layers.size()));
 
         } else if(f.type == "reorg" || f.type == "reorg3d") {
-            netLayers.push_back(new tk::dnn::Reorg(net, f.stride_x));
+            netLayers.push_back(new tk::dnn::Reorg(net, f.stride_x, f.type == "reorg3d"));
 
         } else if(f.type == "region") {
             netLayers.push_back(new tk::dnn::Region(net, f.classes, f.coords, f.num));
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index 2ae909e..0028f9c 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -481,7 +481,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Reorg *l) {
     //std::cout<<"convert Reorg\n";
 
     //std::cout<<"New plugin REORG\n";
-    IPlugin *plugin = new ReorgRT(l->stride);
+    IPlugin *plugin = new ReorgRT(l->stride, l->reorg3d);
     IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
     checkNULL(lRT);
     return lRT;
diff --git a/src/Reorg.cpp b/src/Reorg.cpp
index fe37775..5f349e0 100644
--- a/src/Reorg.cpp
+++ b/src/Reorg.cpp
@@ -5,9 +5,10 @@
 
 namespace tk { namespace dnn {
 
-Reorg::Reorg(Network *net, int stride) : Layer(net) {
+Reorg::Reorg(Network *net, int stride, bool reorg3d) : Layer(net) {
 
     this->stride = stride;
+    this->reorg3d = reorg3d;
     
     output_dim.n = input_dim.n;
     output_dim.c = input_dim.c*stride*stride;
@@ -24,8 +25,10 @@ Reorg::~Reorg() {
 }
 
 dnnType* Reorg::infer(dataDim_t &dim, dnnType* srcData) {
-
-    reorgForward(srcData, dstData, output_dim.n, output_dim.c, output_dim.h, output_dim.w, stride);
+    if (reorg3d)
+        reorgForward(srcData, dstData, output_dim.n, output_dim.c, output_dim.h, output_dim.w, stride);
+    else
+        reorgForward(srcData, dstData, dim.n, dim.c, dim.h, dim.w, stride);
 
     dim = output_dim;
     return dstData;