Merge pull request #428 from kwcckw/dev

Updated ColorPaper.
sparkfish · Mar 30, 2024 · 2a36dba · 2a36dba
2 parents 45c3a7e + 2ca7287
commit 2a36dba
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/augraphy/augmentations/colorpaper.py b/augraphy/augmentations/colorpaper.py
@@ -51,9 +51,13 @@ def add_color(self, image):
 
         # convert to hsv colorspace
         image_hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+
+        random_hue = np.random.randint(self.hue_range[0], self.hue_range[1] + 1)
+        random_saturation = np.random.randint(self.saturation_range[0], self.saturation_range[1] + 1)
+
         # assign hue and saturation
-        image_h = np.random.randint(self.hue_range[0], self.hue_range[1] + 1, size=(ysize, xsize))
-        image_s = np.random.randint(self.saturation_range[0], self.saturation_range[1] + 1, size=(ysize, xsize))
+        image_h = np.random.randint(max(0, random_hue - 5), min(255, random_hue + 5), size=(ysize, xsize))
+        image_s = np.random.randint(max(0, random_saturation - 5), min(255, random_saturation + 5), size=(ysize, xsize))
 
         # assign hue and saturation channel back to hsv image
         image_hsv[:, :, 0] = image_h