Update yolo_training.py

Author: bubbliiiing

nets/yolo_training.py

@@ -1,4 +1,4 @@
-  
+import cv2  
 from random import shuffle
 import numpy as np
 import torch
@@ -349,15 +349,16 @@ def get_random_data(self, annotation_line, input_shape, jitter=.3, hue=.1, sat=1
         hue = rand(-hue, hue)
         sat = rand(1, sat) if rand()<.5 else 1/rand(1, sat)
         val = rand(1, val) if rand()<.5 else 1/rand(1, val)
-        x = rgb_to_hsv(np.array(image)/255.)
-        x[..., 0] += hue
+        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]>360, 0] = 360
+        x[:, :, 1:][x[:, :, 1:]>1] = 1
         x[x<0] = 0
-        image_data = hsv_to_rgb(x)*255 # numpy array, 0 to 1
+        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
 
         # correct boxes
         box_data = np.zeros((len(box),5))
@@ -428,16 +429,17 @@ def get_random_data_with_Mosaic(self, annotation_line, input_shape, hue=.1, sat=
             hue = rand(-hue, hue)
             sat = rand(1, sat) if rand()<.5 else 1/rand(1, sat)
             val = rand(1, val) if rand()<.5 else 1/rand(1, val)
-            x = rgb_to_hsv(np.array(image)/255.)
-            x[..., 0] += hue
+            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]>360, 0] = 360
+            x[:, :, 1:][x[:, :, 1:]>1] = 1
             x[x<0] = 0
-            image = hsv_to_rgb(x)
-
+            image = cv2.cvtColor(x, cv2.COLOR_HSV2RGB) # numpy array, 0 to 1
+            
             image = Image.fromarray((image*255).astype(np.uint8))
             # 将图片进行放置，分别对应四张分割图片的位置
             dx = place_x[index]