perf(data): 测试集进行随机旋转；使用随机像素进行边界填充

Author: zjykzj

rotnet/data/datasets/base_dataset.py

@@ -40,7 +40,7 @@ def __getitem__(self, index):
         if self.target_transform is not None:
             img, target = self.target_transform(img)
         else:
-            # 假定所有训练/测试图像的旋转角度为0
+            # 假定所有训练/测试图像的初始旋转角度为0
             target = 0
 
         # doing this so that it is consistent with all other datasets

rotnet/data/transforms/build.py

@@ -16,7 +16,6 @@
 def build_transform(cfg, train=True):
     size = cfg.MODEL.INPUT_SIZE
 
-    target_transform = None
     if train:
         transform = transforms.Compose([
             transforms.Resize(size),
@@ -26,17 +25,17 @@ def build_transform(cfg, train=True):
             transforms.Normalize((0.5,), (0.5,)),
             transforms.RandomErasing()
         ])
-
-        target_transform = Compose([
-            Rotate(),
-            ToGray(),
-        ])
     else:
         transform = transforms.Compose([
             transforms.Resize(size),
             transforms.Grayscale(),
             transforms.ToTensor(),
             transforms.Normalize((0.5,), (0.5,)),
         ])
+    # 对测试集和训练集都进行随机旋转
+    target_transform = Compose([
+        Rotate(random=True),
+        ToGray(),
+    ])
 
     return transform, target_transform

rotnet/data/transforms/rotate.py

@@ -7,20 +7,39 @@
 @description: 
 """
 
-import math
-import random
-import cv2
 import numpy as np
 
 from rotnet.util.utils import rotate
 
 
 class Rotate:
 
+    def __init__(self, random=False, borderValue=(255, 255, 255)):
+        """
+        :param random: 默认为False，表示使用borderValue指定的边界填充值；如果为True，则忽略borderValue，随机选择填充值
+        :param borderValue: 边界填充值
+        """
+        self.random = random
+        self.borderValue = borderValue
+
     def __call__(self, img: np.ndarray, angle=None):
+        """
+        :param img:
+        :param angle: 如果为None，则随机选择[0,360)的旋转角度
+        :return:
+        """
         assert isinstance(img, np.ndarray)
 
-        angle = random.randint(0, 359)
-        rotate_img = rotate(img, angle)
+        low = 0
+        high = 360
+        if not angle:
+            angle = np.random.randint(low, high=high)
+        if self.random:
+            high = 260
+            borderValue = (
+                np.random.randint(low, high=high), np.random.randint(low, high=high), np.random.randint(low, high=high))
+        else:
+            borderValue = self.borderValue
+        rotate_img = rotate(img, angle, borderValue=borderValue)
 
         return rotate_img, angle

rotnet/engine/trainer.py

@@ -95,7 +95,7 @@ def do_train(cfg, arguments,
             eval_results = do_evaluation(cfg, model, device)
             if summary_writer:
                 for eval_result, dataset_name in zip(eval_results, cfg.DATASETS.TEST):
-                    summary_writer.add_scalar(f'loss/{dataset_name}', eval_result[dataset_name], global_step=iteration)
+                    summary_writer.add_scalar(f'eval/{dataset_name}', eval_result[dataset_name], global_step=iteration)
             model.train()  # *IMPORTANT*: change to train mode after eval.
 
     if summary_writer:

rotnet/util/utils.py

@@ -11,7 +11,7 @@
 import cv2
 
 
-def rotate(img, degree):
+def rotate(img, degree, borderValue=(255, 255, 255)):
     h, w = img.shape[:2]
     center = (w // 2, h // 2)
 
@@ -21,7 +21,7 @@ def rotate(img, degree):
     matrix = cv2.getRotationMatrix2D(center, degree, 1)
     matrix[0, 2] += dst_w // 2 - center[0]
     matrix[1, 2] += dst_h // 2 - center[1]
-    dst_img = cv2.warpAffine(img, matrix, (dst_w, dst_h), borderValue=(255, 255, 255))
+    dst_img = cv2.warpAffine(img, matrix, (dst_w, dst_h), borderValue=borderValue)
 
     # imshow(img, 'src')
     # imshow(dst_img, 'dst')