added augment_equalize_adapthist(img)

yolox/data/data_augment.py

@@ -31,6 +31,14 @@ def augment_hsv(img, hgain=5, sgain=30, vgain=30):
     cv2.cvtColor(img_hsv.astype(img.dtype), cv2.COLOR_HSV2BGR, dst=img)  # no return needed
 
 
+def augment_equalize_adapthist(img):
+    from skimage.exposure import equalize_adapthist
+    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV).astype(np.int16)
+
+    img_hsv[..., 0] = equalize_adapthist(img_hsv[..., 0]).astype(np.int16)
+
+    cv2.cvtColor(img_hsv.astype(img.dtype), cv2.COLOR_HSV2BGR, dst=img)  # no return needed
+
 def get_aug_params(value, center=0):
     if isinstance(value, float):
         return random.uniform(center - value, center + value)
@@ -138,6 +146,11 @@ def _mirror(image, boxes, prob=0.5):
         boxes[:, 0::2] = width - boxes[:, 2::-2]
     return image, boxes
 
+def _equalize_adapthist(image, boxes, prob=0.1):
+    if random.random() < prob:
+        image = augment_equalize_adapthist(image)
+    return image, boxes
+
 
 def preproc(img, input_size, swap=(2, 0, 1)):
     if len(img.shape) == 3:
@@ -182,6 +195,7 @@ def __call__(self, image, targets, input_dim):
 
         if random.random() < self.hsv_prob:
             augment_hsv(image)
+        image, boxes = _equalize_adapthist(image, boxes, prob=0.1)
         image_t, boxes = _mirror(image, boxes, self.flip_prob)
         height, width, _ = image_t.shape
         image_t, r_ = preproc(image_t, input_dim)