Update dataloader.py

utils/dataloader.py

@@ -11,7 +11,7 @@ from torch.utils.data.dataset import Dataset
 from utils.utils import bbox_iou, merge_bboxes
 from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
 from nets.yolo_training import Generator
-
+import cv2
 
 class YoloDataset(Dataset):
     def __init__(self, train_lines, image_size, mosaic=True):
@@ -65,15 +65,16 @@ class YoloDataset(Dataset):
         hue = self.rand(-hue, hue)
         sat = self.rand(1, sat) if self.rand() < .5 else 1 / self.rand(1, sat)
         val = self.rand(1, val) if self.rand() < .5 else 1 / self.rand(1, val)
-        x = rgb_to_hsv(np.array(image) / 255.)
-        x[..., 0] += hue
-        x[..., 0][x[..., 0] > 1] -= 1
-        x[..., 0][x[..., 0] < 0] += 1
+        x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV)
+        x[..., 0] += hue*360
+        x[..., 0][x[..., 0]>1] -= 1
+        x[..., 0][x[..., 0]<0] += 1
         x[..., 1] *= sat
         x[..., 2] *= val
-        x[x > 1] = 1
-        x[x < 0] = 0
-        image_data = hsv_to_rgb(x) * 255  # numpy array, 0 to 1
+        x[x[:,:, 0]>360, 0] = 360
+        x[:, :, 1:][x[:, :, 1:]>1] = 1
+        x[x<0] = 0
+        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
 
         # 调整目标框坐标
         box_data = np.zeros((len(box), 5))
@@ -144,15 +145,16 @@ class YoloDataset(Dataset):
             hue = self.rand(-hue, hue)
             sat = self.rand(1, sat) if self.rand() < .5 else 1 / self.rand(1, sat)
             val = self.rand(1, val) if self.rand() < .5 else 1 / self.rand(1, val)
-            x = rgb_to_hsv(np.array(image) / 255.)
-            x[..., 0] += hue
-            x[..., 0][x[..., 0] > 1] -= 1
-            x[..., 0][x[..., 0] < 0] += 1
+            x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV)
+            x[..., 0] += hue*360
+            x[..., 0][x[..., 0]>1] -= 1
+            x[..., 0][x[..., 0]<0] += 1
             x[..., 1] *= sat
             x[..., 2] *= val
-            x[x > 1] = 1
-            x[x < 0] = 0
-            image = hsv_to_rgb(x)
+            x[x[:,:, 0]>360, 0] = 360
+            x[:, :, 1:][x[:, :, 1:]>1] = 1
+            x[x<0] = 0
+            image = cv2.cvtColor(x, cv2.COLOR_HSV2RGB) # numpy array, 0 to 1
 
             image = Image.fromarray((image * 255).astype(np.uint8))
             # 将图片进行放置，分别对应四张分割图片的位置