refine

Author: ZhenglinZhou

Augmentation.py

@@ -45,7 +45,7 @@ class AutoAugmenter(object):
     Returns:
       A tuple containing the augmented versions of `image` and `bboxes`.
     """
-    def __init__(self, augmentation_name='v2'):
+    def __init__(self, augmentation_name='v4'):
         self.augmentation_name = augmentation_name
 
     def normalizer(self, image, annots):

augmentation_zoo/Myautoaugment_utils.py

@@ -1,8 +1,6 @@
 import numpy as np
 import inspect
 from PIL import Image, ImageOps, ImageEnhance
-import cv2
-import matplotlib.pyplot as plt
 import math
 import random
 
@@ -135,6 +133,16 @@ def policy_v3():
     ]
     return policy
 
+def policy_v4():
+    policy = [
+        [('TranslateX_BBox', 0.6, 4), ('Equalize', 0.8, 10)],
+        [('TranslateY_Only_BBoxes', 0.2, 2), ('Cutout', 0.8, 8)],
+        [('ShearY_BBox', 1.0, 2), ('TranslateY_Only_BBoxes', 0.6, 6)],
+        [('Rotate_BBox', 0.6, 10), ('Color', 1.0, 6)],
+        []
+    ]
+    return policy
+
 """
     AutoContrast
     Equalize
@@ -164,10 +172,11 @@ def policy_v3():
     Cutout_Only_BBoxes
 """
 
+
 def policy_vtest():
     """ (policy, pro, level)"""
     policy = [
-        [('TranslateX_BBox', 1, 10), ('BBox_Cutout', 1, 10)],
+        [('Flip', 1, 10)],
     ]
     return policy
 
@@ -1140,7 +1149,7 @@ def cutout_only_bboxes(image, bboxes, prob, pad_size, replace):
     return _apply_multi_bbox_augmentation_wrapper(
         image, bboxes, prob, cutout, func_changes_bbox, pad_size, replace)
 
-def flip(image, bboxes, prob, level):
+def flip(image, bboxes, prob):
     if np.random.rand() < prob:
         image = image[:, ::-1, :]
 
@@ -1397,7 +1406,7 @@ def distort_image_with_autoaugment(image, bboxes, augmentation_name):
         添加新算法，除函数本身定义外，需在NAME_TO_FUNC、level_to_arg中对应添加。
     """
     available_policies = {'v0': policy_v0, 'v1': policy_v1, 'v2': policy_v2,
-                          'v3': policy_v3, 'test': policy_vtest,
+                          'v3': policy_v3, 'v4': policy_v4, 'test': policy_vtest,
                           'custom': policy_custom}
     if augmentation_name not in available_policies:
         raise ValueError('Invalid augmentation_name: {}'.format(augmentation_name))

augmentation_zoo/mixup.py

@@ -21,7 +21,7 @@
 
 """    AUTOAUGMENT    """
 AUTOAUGMENT = False
-AUTO_POLICY = 'v1'
+AUTO_POLICY = 'v4'
 
 """   GRIDMASK   """
 GRID = False
@@ -38,7 +38,7 @@
 SOA_COPY_TIMES = 3
 SOA_EPOCHS = 30
 SOA_ONE_OBJECT = False
-SOA_ALL_OBJECTS = True
+SOA_ALL_OBJECTS = False
 
 """    MIXUP    """
 MIXUP = False

config.py

@@ -1,7 +1,7 @@
 import torch
 from retinanet import model
 from prepare_data import VocDataset, Normalizer, Resizer, AspectRatioBasedSampler, UnNormalizer, collater
-from Augmentation import autoaugmenter
+from Augmentation import AutoAugmenter
 from torchvision import transforms
 import config
 from retinanet import csv_eval

eval_model.py

@@ -26,13 +26,13 @@ def _make_transform():
     voc_train = VocDataset(VOC_ROOT_DIR, 'train', transform=transforms.Compose(transform_list))
     kitti_train = KittiDataset(KITTI_ROOT_DIR, 'train', transforms.Compose(transform_list))
 
-    for i in range(voc_train.__len__()):
-        print(i)
-        sample = voc_train[i]
-        # easy_visualization(sample)
+    # for i in range(voc_train.__len__()):
+    #     print(i)
+    #     sample = voc_train[i]
+    #     easy_visualization(sample)
 
-    # sample = voc_train[19]
-    # easy_visualization(sample)
+    sample = voc_train[5]
+    easy_visualization(sample)
 
     # for i in range(kitti_train.__len__()):
     #     sample = kitti_train[i]

test_augmentation.py

@@ -135,8 +135,8 @@ def main():
     retinanet.eval()
 
     print('\nBest_mAP:', BEST_MAP_EPOCH)
-    for label in range(retinanet.num_classes()):
-        label_name = retinanet.label_to_name(label)
+    for label in range(dataset_val.num_classes()):
+        label_name = dataset_val.label_to_name(label)
         print('{}: {}'.format(label_name, best_average_precisions[label][0]))
     print('BEST MAP: ', BEST_MAP)
     # torch.save(retinanet, 'model_final.pt')