diff --git a/nets/ssd_training.py b/nets/ssd_training.py index e49b82cf705720f5441663b711326dd1028bcfde..ac9c2ce2963998ad0306412db542d4ec6e80ef55 100644 --- a/nets/ssd_training.py +++ b/nets/ssd_training.py @@ -8,7 +8,7 @@ from utils.config import Config from utils.box_utils import match, log_sum_exp from matplotlib.colors import rgb_to_hsv, hsv_to_rgb from PIL import Image - +import cv2 MEANS = (104, 117, 123) @@ -149,15 +149,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))