Experimenting with random erasing changes

timm/data/random_erasing.py

@@ -43,6 +43,7 @@ class RandomErasing:
         self.sl = sl
         self.sh = sh
         self.min_aspect = min_aspect
+        self.max_count = 8
         mode = mode.lower()
         self.rand_color = False
         self.per_pixel = False
@@ -58,18 +59,20 @@ class RandomErasing:
         if random.random() > self.probability:
             return
         area = img_h * img_w
-        for attempt in range(100):
+        count = random.randint(1, self.max_count)
-            target_area = random.uniform(self.sl, self.sh) * area
+        for _ in range(count):
-            aspect_ratio = random.uniform(self.min_aspect, 1 / self.min_aspect)
+            for attempt in range(10):
-            h = int(round(math.sqrt(target_area * aspect_ratio)))
+                target_area = random.uniform(self.sl / count, self.sh / count) * area
-            w = int(round(math.sqrt(target_area / aspect_ratio)))
+                aspect_ratio = random.uniform(self.min_aspect, 1 / self.min_aspect)
-            if w < img_w and h < img_h:
+                h = int(round(math.sqrt(target_area * aspect_ratio)))
-                top = random.randint(0, img_h - h)
+                w = int(round(math.sqrt(target_area / aspect_ratio)))
-                left = random.randint(0, img_w - w)
+                if w < img_w and h < img_h:
-                img[:, top:top + h, left:left + w] = _get_pixels(
+                    top = random.randint(0, img_h - h)
-                    self.per_pixel, self.rand_color, (chan, h, w),
+                    left = random.randint(0, img_w - w)
-                    dtype=dtype, device=self.device)
+                    img[:, top:top + h, left:left + w] = _get_pixels(
-                break
+                        self.per_pixel, self.rand_color, (chan, h, w),
+                        dtype=dtype, device=self.device)
+                    break
 
     def __call__(self, input):
         if len(input.size()) == 3: