yjh0410 1 year ago
parent
commit
9e1a7cea0a
2 changed files with 11 additions and 7 deletions
  1. 5 2
      dataset/data_augment/ssd_augment.py
  2. 6 5
      dataset/data_augment/yolov5_augment.py

+ 5 - 2
dataset/data_augment/ssd_augment.py

@@ -54,10 +54,13 @@ class Resize(object):
 
     def __call__(self, image, boxes=None, labels=None):
         orig_h, orig_w = image.shape[:2]
-        image = cv2.resize(image, (self.img_size, self.img_size))
+
+        # resize
+        image = cv2.resize(image, (self.img_size, self.img_size)).astype(np.float32)
+        img_h, img_w = image.shape[:2]
+
         # rescale bbox
         if boxes is not None:
-            img_h, img_w = image.shape[:2]
             boxes[..., [0, 2]] = boxes[..., [0, 2]] / orig_w * img_w
             boxes[..., [1, 3]] = boxes[..., [1, 3]] / orig_h * img_h
 

+ 6 - 5
dataset/data_augment/yolov5_augment.py

@@ -146,6 +146,12 @@ class YOLOv5Augmentation(object):
             img = image
         img_h, img_w = img.shape[:2]
 
+        # rescale bbox
+        boxes_ = target["boxes"].copy()
+        boxes_[:, [0, 2]] = boxes_[:, [0, 2]] / img_w0 * img_w
+        boxes_[:, [1, 3]] = boxes_[:, [1, 3]] / img_h0 * img_h
+        target["boxes"] = boxes_
+        
         # --------------- Filter bad targets ---------------
         tgt_boxes_wh = target["boxes"][..., 2:] - target["boxes"][..., :2]
         min_tgt_size = np.min(tgt_boxes_wh, axis=-1)
@@ -166,11 +172,6 @@ class YOLOv5Augmentation(object):
         # --------------- Spatial augmentations ---------------
         ## Random perspective
         if not mosaic:
-            # rescale bbox
-            boxes_ = target["boxes"].copy()
-            boxes_[:, [0, 2]] = boxes_[:, [0, 2]] / img_w0 * img_w
-            boxes_[:, [1, 3]] = boxes_[:, [1, 3]] / img_h0 * img_h
-            target["boxes"] = boxes_
             # spatial augment
             target_ = np.concatenate(
                 (target['labels'][..., None], target['boxes']), axis=-1)