fix: allows to specify size and color/bw with segmentation models

jolibrain · Sep 15, 2020 · 58ecb4a · 58ecb4a
1 parent 53edf52
commit 58ecb4a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 0 deletions.
diff --git a/src/backends/caffe/caffelib.cc b/src/backends/caffe/caffelib.cc
@@ -295,6 +295,10 @@ namespace dd
           {
             update_protofile_imageDataLayer(net_param);
           }
+        else if (this->_inputc._segmentation)
+          {
+            update_protofile_denseImageDataLayer(net_param);
+          }
 
         // input should be ok, now do the output
         if (this->_inputc._multi_label)
@@ -638,6 +642,8 @@ namespace dd
                 this->_inputc.height());
             lparam->mutable_dense_image_data_param()->set_new_width(
                 this->_inputc.width());
+            lparam->mutable_dense_image_data_param()->set_is_color(
+                this->_inputc.channels() == 3);
             // XXX: DenseImageData supports crop_height and crop_width
           }
       }
@@ -4882,6 +4888,21 @@ namespace dd
     lparam->clear_transform_param();
   }
 
+  template <class TInputConnectorStrategy, class TOutputConnectorStrategy,
+            class TMLModel>
+  void CaffeLib<TInputConnectorStrategy, TOutputConnectorStrategy, TMLModel>::
+      update_protofile_denseImageDataLayer(caffe::NetParameter &net_param)
+  {
+    caffe::LayerParameter *lparam = net_param.mutable_layer(0);
+    caffe::DenseImageDataParameter *image_data_parameter
+        = lparam->mutable_dense_image_data_param();
+    image_data_parameter->set_new_height(this->_inputc.height());
+    image_data_parameter->set_new_width(this->_inputc.width());
+    ImgCaffeInputFileConn *timg
+        = reinterpret_cast<ImgCaffeInputFileConn *>(&this->_inputc);
+    image_data_parameter->set_is_color(!timg->_bw);
+  }
+
   template <class TInputConnectorStrategy, class TOutputConnectorStrategy,
             class TMLModel>
   bool CaffeLib<TInputConnectorStrategy, TOutputConnectorStrategy,

diff --git a/src/backends/caffe/caffelib.h b/src/backends/caffe/caffelib.h
@@ -246,6 +246,8 @@ namespace dd
 
     void update_protofile_imageDataLayer(caffe::NetParameter &net_param);
 
+    void update_protofile_denseImageDataLayer(caffe::NetParameter &net_param);
+
     void update_protofiles_dice_params(caffe::DiceCoefLossParameter *dclp,
                                        const APIData &ad);