fix: dataset output dimension after crop augmentation

jolibrain · Jan 8, 2022 · 636d455 · 636d455
1 parent ac7ce0f
commit 636d455
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 15 deletions.
diff --git a/src/backends/torch/torchdataset.cc b/src/backends/torch/torchdataset.cc
@@ -242,18 +242,10 @@ namespace dd
 
         if (_segmentation)
           {
-
             cv::resize(bw_target, bw_target, cv::Size(width, height), 0, 0,
                        cv::INTER_NEAREST);
           }
       }
-
-    if (_segmentation)
-      {
-        at::Tensor targett_seg
-            = image_to_tensor(bw_target, height, width, true);
-        targett.push_back(targett_seg);
-      }
   }
 
   // add image batch
@@ -554,14 +546,13 @@ namespace dd
                       _img_rand_aug_cv.augment(bgr);
                   }
 
-                torch::Tensor imgt
-                    = image_to_tensor(bgr, inputc->height(), inputc->width());
+                torch::Tensor imgt = image_to_tensor(bgr, bgr.rows, bgr.cols);
                 d.push_back(imgt);
 
                 if (_segmentation)
                   {
                     at::Tensor targett_seg = image_to_tensor(
-                        bw_target, inputc->height(), inputc->width(), true);
+                        bw_target, bw_target.rows, bw_target.cols, true);
                     t.push_back(targett_seg);
                   }
               }
@@ -834,8 +825,6 @@ namespace dd
     return imgt;
   }
 
-  // TODO: segmentation target image to tensor
-
   at::Tensor TorchDataset::target_to_tensor(const int &target)
   {
     at::Tensor targett{ torch::full(1, target, torch::kLong) };

diff --git a/src/backends/torch/torchlib.cc b/src/backends/torch/torchlib.cc
@@ -620,7 +620,7 @@ namespace dd
         throw;
       }
 
-    // TODO: set inputc dataset data augmentation options
+    // set inputc dataset data augmentation options
     APIData ad_mllib = ad.getobj("parameters").getobj("mllib");
     if (typeid(inputc) == typeid(ImgTorchInputFileConn))
       {

diff --git a/src/backends/torch/torchloss.cc b/src/backends/torch/torchloss.cc
@@ -121,7 +121,6 @@ namespace dd
       {
         if (_loss.empty())
           {
-
             loss = torch::nn::functional::cross_entropy(
                 y_pred, y.squeeze(1).to(torch::kLong)); // TODO: options
           }