refine small_object_augmentation

Author: ZhenglinZhou

augmentation_zoo/MyGridMask.py

@@ -12,6 +12,7 @@ def __init__(self, use_h, use_w, rotate=1, offset=False, ratio=0.5, mode=0, prob
         self.mode = mode
         self.st_prob = prob
         self.prob = prob
+
     def set_prob(self, epoch, max_epoch):
         self.prob = self.st_prob * epoch / max_epoch
 

augmentation_zoo/SmallObjectAugmentation.py

@@ -2,16 +2,19 @@
 import random
 
 class SmallObjectAugmentation(object):
-    def __init__(self, thresh=64*64, prob=0, copy_times=3, all_objects=False, one_object=False):
+    def __init__(self, thresh=64*64, prob=0.5, copy_times=3, epochs=30, all_objects=False, one_object=False):
         """
+        sample = {'img':img, 'annot':annots}
         img = [height, width, 3]
         annot = [xmin, ymin, xmax, ymax, label]
-        thresh：小目标边界
-
+        thresh：the detection threshold of the small object. If annot_h * annot_w < thresh, the object is small
+        prob: the prob to do small object augmentation
+        epochs: the epochs to do
         """
         self.thresh = thresh
         self.prob = prob
         self.copy_times = copy_times
+        self.epochs = epochs
         self.all_objects = all_objects
         self.one_object = one_object
         if self.all_objects or self.one_object:
@@ -42,27 +45,22 @@ def donot_overlap(self, new_annot, annots):
 
     def create_copy_annot(self, h, w, annot, annots):
         annot = annot.astype(np.int)
-        new_annot = list()
         annot_h, annot_w = annot[3] - annot[1], annot[2] - annot[0]
-        random_x, random_y = np.random.randint(int(annot_w/2), int(w-annot_w/2)), \
-                             np.random.randint(int(annot_h/2), int(h-annot_h/2))
-
-        if np.int(random_x - annot_w/2) < 0 or np.floor(random_x + annot_w/2) > w or \
-                np.int(random_y - annot_h/2) < 0 or np.floor(random_y + annot_h/2) > h:
-            return self.create_copy_annot(h ,w, annot, annots)
-
-        xmin, ymin = random_x - annot_w/2, random_y - annot_h/2
-        xmax, ymax = xmin + annot_w, ymin + annot_h
-        new_annot.append(xmin), new_annot.append(ymin)
-        new_annot.append(xmax), new_annot.append(ymax)
-        new_annot.append(annot[4])
-
-        new_annot = np.array(new_annot).astype(np.int)
-
-        if self.donot_overlap(new_annot, annots) is False:
-            return self.create_copy_annot(h, w, annot, annots)
-
-        return new_annot
+        for epoch in range(self.epochs):
+            random_x, random_y = np.random.randint(int(annot_w / 2), int(w - annot_w / 2)), \
+                                 np.random.randint(int(annot_h / 2), int(h - annot_h / 2))
+            xmin, ymin = random_x - annot_w / 2, random_y - annot_h / 2
+            xmax, ymax = xmin + annot_w, ymin + annot_h
+            if np.int(xmin) < 0 or np.floor(xmax) > w or \
+                    np.int(ymin) < 0 or np.floor(ymax) > h:
+                continue
+            new_annot = np.array([xmin, ymin, xmax, ymax, annot[4]]).astype(np.int)
+
+            if self.donot_overlap(new_annot, annots) is False:
+                continue
+
+            return new_annot
+        return None
 
     def add_patch_in_img(self, annot, copy_annot, image):
         copy_annot = copy_annot.astype(np.int)
@@ -74,15 +72,32 @@ def __call__(self, sample):
         if np.random.rand() > self.prob: return sample
 
         img, annots = sample['img'], sample['annot']
-        h, w, l = img.shape[0], img.shape[1], annots.shape[0]
+        h, w= img.shape[0], img.shape[1]
+
+        small_object_list = list()
+        for idx in range(annots.shape[0]):
+            annot = annots[idx]
+            annot_h, annot_w = annot[3] - annot[1], annot[2] - annot[0]
+            if self.issmallobject(annot_h, annot_w):
+                small_object_list.append(idx)
 
+        l = len(small_object_list)
+        # No Small Object
+        if l == 0: return sample
+
+        # Refine the copy_object by the given policy
+        # Policy 2:
         copy_object_num = np.random.randint(0, l)
+        # Policy 3:
         if self.all_objects:
             copy_object_num = l
+        # Policy 1:
         if self.one_object:
             copy_object_num = 1
+
         random_list = random.sample(range(l), copy_object_num)
-        select_annots = annots[random_list, :]
+        annot_idx_of_small_object = [small_object_list[idx] for idx in random_list]
+        select_annots = annots[annot_idx_of_small_object, :]
         annots = annots.tolist()
         for idx in range(copy_object_num):
             annot = select_annots[idx]
@@ -91,7 +106,7 @@ def __call__(self, sample):
             if self.issmallobject(annot_h, annot_w) is False: continue
 
             for i in range(self.copy_times):
-                new_annot = self.create_copy_annot(h, w, annot, annots)
+                new_annot = self.create_copy_annot(h, w, annot, annots,)
                 if new_annot is not None:
                     img = self.add_patch_in_img(new_annot, annot, img)
                     annots.append(new_annot)

config.py

@@ -36,8 +36,9 @@
 SOA_THRESH = 64*64
 SOA_PROB = 1
 SOA_COPY_TIMES = 3
+SOA_EPOCHS = 30
 SOA_ONE_OBJECT = False
-SOA_ALL_OBJECTS = False
+SOA_ALL_OBJECTS = True
 
 """    MIXUP    """
 MIXUP = False

test_augmentation.py

@@ -18,7 +18,7 @@ def _make_transform():
     if cfg.RANDOM_FLIP:
         transform_list.append(RandomFlip())
     if cfg.SMALL_OBJECT_AUGMENTATION:
-        transform_list.append(SmallObjectAugmentation(cfg.SOA_THRESH, cfg.SOA_PROB, cfg.SOA_COPY_TIMES, cfg.SOA_ALL_OBJECTS, cfg.SOA_ONE_OBJECT))
+        transform_list.append(SmallObjectAugmentation(cfg.SOA_THRESH, cfg.SOA_PROB, cfg.SOA_COPY_TIMES, cfg.SOA_EPOCHS, cfg.SOA_ALL_OBJECTS, cfg.SOA_ONE_OBJECT))
     return transform_list
 
 if __name__ == '__main__':

train.py

@@ -31,7 +31,7 @@ def _make_transform():
     if cfg.RANDOM_FLIP:
         transform_list.append(RandomFlip())
     if cfg.SMALL_OBJECT_AUGMENTATION:
-        transform_list.append(SmallObjectAugmentation(cfg.SOA_THRESH, cfg.SOA_PROB, cfg.SOA_COPY_TIMES, cfg.SOA_ALL_OBJECTS, cfg.SOA_ONE_OBJECT))
+        transform_list.append(SmallObjectAugmentation(cfg.SOA_THRESH, cfg.SOA_PROB, cfg.SOA_COPY_TIMES, cfg.SOA_EPOCHS, cfg.SOA_ALL_OBJECTS, cfg.SOA_ONE_OBJECT))
     transform_list.append(Normalizer())
     transform_list.append(Resizer())
     return transform_list