diff --git a/nets/yolo_training.py b/nets/yolo_training.py index 852f73837ea5312893330eeaa67fcc2eb48adabc..0935681d0a65de0fa89d104b8b85958d09db9302 100644 --- a/nets/yolo_training.py +++ b/nets/yolo_training.py @@ -1,4 +1,4 @@ - +import cv2 from random import shuffle import numpy as np import torch @@ -349,15 +349,16 @@ class Generator(object): 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 @@ class Generator(object): 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]