added a force_encoded_color flag to the data layer. Printing a warning if images of different channel dimensions are encoded together

1 files changed, 13 insertions, 2 deletions

diff --git a/src/caffe/layers/data_layer.cpp b/src/caffe/layers/data_layer.cpp
index 891d0391..7716406f 100644
--- a/src/caffe/layers/data_layer.cpp
+++ b/src/caffe/layers/data_layer.cpp
@@ -42,7 +42,9 @@ void DataLayer<Dtype>::DataLayerSetUp(const vector<Blob<Dtype>*>& bottom,
   Datum datum;
   datum.ParseFromString(cursor_->value());
 
-  if (DecodeDatumNative(&datum)) {
+  bool force_color = this->layer_param_.data_param().force_encoded_color();
+  if ((force_color && DecodeDatum(&datum, true)) ||
+      DecodeDatumNative(&datum)) {
     LOG(INFO) << "Decoding Datum";
   }
   // image
@@ -90,6 +92,7 @@ void DataLayer<Dtype>::InternalThreadEntry() {
     top_label = this->prefetch_label_.mutable_cpu_data();
   }
   const int batch_size = this->layer_param_.data_param().batch_size();
+  bool force_color = this->layer_param_.data_param().force_encoded_color();
   for (int item_id = 0; item_id < batch_size; ++item_id) {
     timer.Start();
     // get a blob
@@ -98,7 +101,15 @@ void DataLayer<Dtype>::InternalThreadEntry() {
 
     cv::Mat cv_img;
     if (datum.encoded()) {
-       cv_img = DecodeDatumToCVMatNative(datum);
+      if (force_color)
+        cv_img = DecodeDatumToCVMat(datum, true);
+      else
+        cv_img = DecodeDatumToCVMatNative(datum);
+      if (cv_img.channels() != this->transformed_data_.channels())
+        LOG(WARNING) << "Your dataset contains encoded images with mixed "
+        << "channel sizes. Consider adding a 'force_color' flag to the "
+        << "model definition, or rebuild your dataset using "
+        << "convert_imageset.";
     }
     read_time += timer.MicroSeconds();
     timer.Start();