fix(torch): data augmentation handle dummy bboxes correctly

jolibrain · Jan 31, 2023 · 53d0c39 · 53d0c39
1 parent 041b649
commit 53d0c39
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/src/backends/torch/torchdataset.cc b/src/backends/torch/torchdataset.cc
@@ -395,12 +395,24 @@ namespace dd
         std::vector<torch::Tensor> t_sample = t;
         if (_segmentation)
           bw_target_sample = bw_target.clone();
+        bool no_bbox = false;
+        if (_bbox)
+          {
+            auto &cls_bbox = t_sample.at(1);
+            no_bbox = cls_bbox.numel() == 1 && cls_bbox.item<int>() == 0;
+          }
 
         // data augmentation can apply here, with OpenCV
         if (!_test)
           {
             if (_bbox)
-              _img_rand_aug_cv.augment_with_bbox(bgr_sample, t_sample);
+              {
+                // don't augment dummy bbox
+                if (no_bbox)
+                  _img_rand_aug_cv.augment(bgr_sample);
+                else
+                  _img_rand_aug_cv.augment_with_bbox(bgr_sample, t_sample);
+              }
             else if (_segmentation)
               _img_rand_aug_cv.augment_with_segmap(bgr_sample,
                                                    bw_target_sample);
@@ -412,7 +424,12 @@ namespace dd
             // cropping requires test set 'augmentation'
             if (_bbox)
               {
-                _img_rand_aug_cv.augment_test_with_bbox(bgr_sample, t_sample);
+                // don't augment dummy bbox
+                if (no_bbox)
+                  _img_rand_aug_cv.augment_test(bgr_sample);
+                else
+                  _img_rand_aug_cv.augment_test_with_bbox(bgr_sample,
+                                                          t_sample);
               }
             else if (_segmentation)
               _img_rand_aug_cv.augment_test_with_segmap(bgr_sample,

diff --git a/tools/torch/trace_yolox.py b/tools/torch/trace_yolox.py
@@ -25,7 +25,7 @@ def main():
 - Export dd-trained yolox model to onnx for trt inference:
 
     python3 trace_yolox.py -v yolox[-s|-m|-l|...] --yolox_path [YOLOX_PATH] --output_dir [OUTPUT_DIR] --from_repo [DD_REPO] --to_onnx
-"""
+""", formatter_class=argparse.RawDescriptionHelpFormatter
     )
     parser.add_argument("model", type=str, help="Model to export")
     parser.add_argument(