Fix random bounding box crop

dali/pipeline/operators/crop/bbox_crop.cc

@@ -23,13 +23,15 @@ two Tensors: `BBoxes` containing bounding boxes represented as `[l,t,r,b]` or `[
 corresponding label for each bounding box. Resulting prospective crop is provided as two Tensors: `Begin` containing the starting
 coordinates for the `crop` in `(x,y)` format, and 'Size' containing the dimensions of the `crop` in `(w,h)` format.
 Bounding boxes are provided as a `(m*4)` Tensor, where each bounding box is represented as `[l,t,r,b]` or `[x,y,w,h]`. Resulting
-labels match the boxes that remain, after being discarded with respect to the minimum accepted intersection threshold.)code")
+labels match the boxes that remain, after being discarded with respect to the minimum accepted intersection threshold.
+Be advised, when `allow_no_crop` is `false` and `thresholds` does not contain `0` it is good to increase `num_attempts` as otherwise
+it may loop for a very long time.)code")
     .NumInput(2)
     .NumOutput(4)
     .AddOptionalArg(
         "thresholds",
         R"code(Minimum overlap (Intersection over union) of the bounding boxes with respect to the prospective crop.
-Selected at random for every sample from provided values. Default leaves the input image as-is in the new crop.)code",
+Selected at random for every sample from provided values. Default imposes no restrictions on Intersection over Union for boxes and crop.)code",
         std::vector<float>{0.f})
     .AddOptionalArg(
         "aspect_ratio",
@@ -48,6 +50,10 @@ Default values disallow changes in aspect ratio.)code",
         "num_attempts",
         R"code(Number of attempts to retrieve a patch with the desired parameters.)code",
         1)
+    .AddOptionalArg(
+        "allow_no_crop",
+        R"code(It true, includes no cropping as one of the random options.)code",
+        true)
     .EnforceInputLayout(DALI_NHWC);
 
 template <>
@@ -93,7 +99,7 @@ void RandomBBoxCrop<CPUBackend>::WriteBoxesToOutput(
 
 template <>
 void RandomBBoxCrop<CPUBackend>::WriteLabelsToOutput(
-    SampleWorkspace *ws, const std::vector<int> &labels) {
+  SampleWorkspace *ws, const std::vector<int> &labels) {
   auto &labels_out = ws->Output<CPUBackend>(3);
   labels_out.Resize({static_cast<Index>(labels.size()), 1});
 
@@ -129,18 +135,20 @@ void RandomBBoxCrop<CPUBackend>::RunImpl(SampleWorkspace *ws, const int) {
     labels.emplace_back(*label_data);
   }
 
-  const auto prospective_crop =
-      FindProspectiveCrop(bounding_boxes, labels, SelectMinimumOverlap());
+  ProspectiveCrop prospective_crop;
+  while (!prospective_crop.success)
+    prospective_crop  = FindProspectiveCrop(
+        bounding_boxes, labels, SelectMinimumOverlap());
 
-  const auto &selected_boxes = std::get<1>(prospective_crop);
-  const auto &selected_labels = std::get<2>(prospective_crop);
+  const auto &selected_boxes = prospective_crop.boxes;
+  const auto &selected_labels = prospective_crop.labels;
 
   DALI_ENFORCE(selected_boxes.size() == selected_labels.size(),
-               "Expected boxes.size() == labels.size(). Received: " +
-                   std::to_string(selected_boxes.size()) +
-                   "!=" + std::to_string(selected_labels.size()));
+              "Expected boxes.size() == labels.size(). Received: " +
+                  std::to_string(selected_boxes.size()) +
+                  "!=" + std::to_string(selected_labels.size()));
 
-  WriteCropToOutput(ws, std::get<0>(prospective_crop));
+  WriteCropToOutput(ws, prospective_crop.crop);
   WriteBoxesToOutput(ws, selected_boxes);
   WriteLabelsToOutput(ws, selected_labels);
 }

dali/pipeline/operators/crop/bbox_crop.h

@@ -50,21 +50,34 @@ class RandomBBoxCrop : public Operator<Backend> {
   using Crop = BoundingBox;
   using BoundingBoxes = std::vector<BoundingBox>;
 
+  struct ProspectiveCrop {
+    bool success = false;
+    Crop crop = Crop::FromLtrb(0, 0, 1, 1);
+    BoundingBoxes boxes;
+    std::vector<int> labels;
+
+    ProspectiveCrop(
+      bool success, const Crop &crop, const BoundingBoxes &boxes, const std::vector<int> &labels) :
+      success(success), crop(crop), boxes(boxes), labels(labels) {}
+    ProspectiveCrop() = default;
+  };
+
  public:
   explicit inline RandomBBoxCrop(const OpSpec &spec)
       : Operator<Backend>(spec),
-        thresholds_{spec.GetRepeatedArgument<float>("thresholds")},
         scaling_bounds_{Bounds(spec.GetRepeatedArgument<float>("scaling"))},
         aspect_ratio_bounds_{
             Bounds(spec.GetRepeatedArgument<float>("aspect_ratio"))},
         ltrb_{spec.GetArgument<bool>("ltrb")},
         num_attempts_{spec.GetArgument<int>("num_attempts")}
 
   {
-    DALI_ENFORCE(!thresholds_.empty(),
+    auto thresholds = spec.GetRepeatedArgument<float>("thresholds");
+
+    DALI_ENFORCE(!thresholds.empty(),
                  "At least one threshold value must be provided");
 
-    for (const auto &threshold : thresholds_) {
+    for (const auto &threshold : thresholds) {
       DALI_ENFORCE(0.0 <= threshold,
                    "Threshold value must be >= 0.0. Received: " +
                        std::to_string(threshold));
@@ -73,7 +86,12 @@ class RandomBBoxCrop : public Operator<Backend> {
                        std::to_string(threshold));
       DALI_ENFORCE(num_attempts_ > 0,
                    "Minimum number of attempts must be greater than zero");
+
+      sample_options_.push_back(std::make_pair(threshold, true));
     }
+
+    if (spec.GetArgument<bool>("allow_no_crop"))
+      sample_options_.push_back(std::make_pair(0.f, false));
   }
 
   ~RandomBBoxCrop() override = default;
@@ -88,24 +106,22 @@ class RandomBBoxCrop : public Operator<Backend> {
 
   void WriteLabelsToOutput(SampleWorkspace *ws, const std::vector<int> &labels);
 
-  float SelectMinimumOverlap() {
-    static std::uniform_int_distribution<> sampler(
-        0, static_cast<int>(thresholds_.size() - 1));
-    return thresholds_[sampler(rd_)];
+  std::pair<float, bool> SelectMinimumOverlap() {
+    std::uniform_int_distribution<> sampler(
+      0, static_cast<int>(sample_options_.size() - 1));
+    return sample_options_[sampler(rd_)];
   }
 
   float SampleCandidateDimension() {
-    static std::uniform_real_distribution<> sampler(scaling_bounds_.min,
-                                                    scaling_bounds_.max);
+    std::uniform_real_distribution<> sampler(scaling_bounds_.min, scaling_bounds_.max);
     return static_cast<float>(sampler(rd_));
   }
 
   bool ValidAspectRatio(float width, float height) const {
     return aspect_ratio_bounds_.Contains(width / height);
   }
 
-  bool ValidOverlap(const Crop &crop, const BoundingBoxes &boxes,
-                    float threshold) {
+  bool ValidOverlap(const Crop &crop, const BoundingBoxes &boxes, float threshold) {
     return std::all_of(boxes.begin(), boxes.end(),
                        [&crop, threshold](const BoundingBox &box) {
                          return crop.IntersectionOverUnion(box) >= threshold;
@@ -164,42 +180,41 @@ class RandomBBoxCrop : public Operator<Backend> {
     return std::make_pair(candidate_boxes, candidate_labels);
   }
 
-  std::tuple<Crop, BoundingBoxes, std::vector<int>> FindProspectiveCrop(
+  ProspectiveCrop FindProspectiveCrop(
       const BoundingBoxes &bounding_boxes, const std::vector<int> &labels,
-      float minimum_overlap) {
-    if (minimum_overlap > 0) {
-      for (int i = 0; i < num_attempts_; ++i) {
-        // Image is HWC
-        const auto candidate_height = SampleCandidateDimension();
-        const auto candidate_width = SampleCandidateDimension();
+      std::pair<float, bool> minimum_overlap) {
+    if (!minimum_overlap.second)
+      return ProspectiveCrop(true, Crop::FromLtrb(0, 0, 1, 1), bounding_boxes, labels);
+
+    for (int i = 0; i < num_attempts_; ++i) {
+      // Image is HWC
+      const auto candidate_height = SampleCandidateDimension();
+      const auto candidate_width = SampleCandidateDimension();
 
-        if (ValidAspectRatio(candidate_height, candidate_width)) {
-          const auto candidate_crop =
-              SamplePatch(candidate_height, candidate_width);
+      if (ValidAspectRatio(candidate_height, candidate_width)) {
+        const auto candidate_crop =
+            SamplePatch(candidate_height, candidate_width);
 
+        if (ValidOverlap(candidate_crop, bounding_boxes, minimum_overlap.first)) {
           BoundingBoxes candidate_boxes;
           std::vector<int> candidate_labels;
 
           std::tie(candidate_boxes, candidate_labels) =
-              DiscardBoundingBoxesByCentroid(candidate_crop, bounding_boxes,
-                                             labels);
+                DiscardBoundingBoxesByCentroid(candidate_crop, bounding_boxes, labels);
 
-          if (ValidOverlap(candidate_crop, candidate_boxes, minimum_overlap)) {
-            const auto remapped_boxes =
-                RemapBoxes(candidate_crop, candidate_boxes, candidate_height,
-                           candidate_width);
-
-            return std::make_tuple(candidate_crop, remapped_boxes,
-                                   candidate_labels);
+          if (candidate_boxes.begin() != candidate_boxes.end()) {
+            const auto remapped_boxes = RemapBoxes(
+              candidate_crop, candidate_boxes, candidate_height, candidate_width);
+            return ProspectiveCrop(true, candidate_crop, remapped_boxes, candidate_labels);
           }
         }
       }
     }
-
-    return std::make_tuple(Crop::FromLtrb(0, 0, 1, 1), bounding_boxes, labels);
+    return ProspectiveCrop();
   }
 
-  const std::vector<float> thresholds_;
+  // float - threshold for IoU, bool - whether to apply crop (false means no cropp)
+  std::vector<std::pair<float, bool>> sample_options_;
   const Bounds scaling_bounds_;
   const Bounds aspect_ratio_bounds_;
   const bool ltrb_;