From e4c866e3343048bb5ab620ee0d61c7ed76bf16be Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=99=BD=E9=B9=AD=E5=85=88=E7=94=9F?= <766529835@qq.com>
Date: Sat, 4 Feb 2023 18:59:37 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E9=A9=AC=E8=B5=9B=E5=85=8B?=
 =?UTF-8?q?=E6=95=B0=E6=8D=AE=E5=A2=9E=E5=BC=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 utils/dataloader.py | 34 +++++++++++++++++-----------------
 1 file changed, 17 insertions(+), 17 deletions(-)

diff --git a/utils/dataloader.py b/utils/dataloader.py
index 4c40d5b..5dcdf88 100644
--- a/utils/dataloader.py
+++ b/utils/dataloader.py
@@ -41,18 +41,18 @@ class YoloDataset(Dataset):
         #   训练时进行数据的随机增强
         #   验证时不进行数据的随机增强
         #---------------------------------------------------#
-        # if self.mosaic and self.rand() < self.mosaic_prob and self.epoch_now < self.epoch_length * self.special_aug_ratio:
-        lines = sample(self.annotation_lines, 3)
-        lines.append(self.annotation_lines[index])
-        shuffle(lines)
-        image, rbox  = self.get_random_data_with_Mosaic(lines, self.input_shape)
-            
-        #     if self.mixup and self.rand() < self.mixup_prob:
-        #         lines           = sample(self.annotation_lines, 1)
-        #         image_2, rbox_2  = self.get_random_data(lines[0], self.input_shape, random = self.train)
-        #         image, rbox      = self.get_random_data_with_MixUp(image, rbox, image_2, rbox_2)
-        # else:
-        #     image, rbox      = self.get_random_data(self.annotation_lines[index], self.input_shape, random = self.train)
+        if self.mosaic and self.rand() < self.mosaic_prob and self.epoch_now < self.epoch_length * self.special_aug_ratio:
+            lines = sample(self.annotation_lines, 3)
+            lines.append(self.annotation_lines[index])
+            shuffle(lines)
+            image, rbox  = self.get_random_data_with_Mosaic(lines, self.input_shape)
+                
+            if self.mixup and self.rand() < self.mixup_prob:
+                lines           = sample(self.annotation_lines, 1)
+                image_2, rbox_2  = self.get_random_data(lines[0], self.input_shape, random = self.train)
+                image, rbox      = self.get_random_data_with_MixUp(image, rbox, image_2, rbox_2)
+        else:
+            image, rbox      = self.get_random_data(self.annotation_lines[index], self.input_shape, random = self.train)
 
         image       = np.transpose(preprocess_input(np.array(image, dtype=np.float32)), (2, 0, 1))
         rbox        = np.array(rbox, dtype=np.float32)
@@ -186,11 +186,11 @@ class YoloDataset(Dataset):
                 rbox[:, 0] = 1 - rbox[:, 0]
                 rbox[:, 4] *= -1
         # 查看旋转框是否正确
-        draw = ImageDraw.Draw(image)
-        polys = rbox2poly(rbox[..., :5])*w
-        for poly in polys:
-            draw.polygon(xy=list(poly))
-        image.show()
+        # draw = ImageDraw.Draw(image)
+        # polys = rbox2poly(rbox[..., :5])*w
+        # for poly in polys:
+        #     draw.polygon(xy=list(poly))
+        # image.show()
         return image_data, rbox
     
     def merge_rboxes(self, rboxes, cutx, cuty):
-- 
GitLab