[WIP] Add RandAugment

ffcv/libffcv.py

@@ -25,6 +25,9 @@ def read(fileno:int, destination:np.ndarray, offset:int):
 ctypes_add_weighted = lib.add_weighted
 ctypes_add_weighted.argtypes = [c_int64, c_float, c_int64, c_float, c_int64, c_int64, c_int64]
 
+ctypes_equalize = lib.equalize
+ctypes_equalize.argtypes = [c_int64, c_int64, c_int64, c_int64]
+
 
 def resize_crop(source, start_row, end_row, start_col, end_col, destination):
     ctypes_resize(0,

ffcv/transforms/utils/fast_crop.py

@@ -1,7 +1,58 @@
 import ctypes
-from numba import njit
+from numba import njit, prange
 import numpy as np
-from ...libffcv import ctypes_resize, ctypes_rotate, ctypes_shear, ctypes_add_weighted
+from ...libffcv import ctypes_resize, ctypes_rotate, ctypes_shear, ctypes_add_weighted, ctypes_equalize
+
+
+"""
+Custom equalize -- equivalent to torchvision.transforms.functional.equalize,
+but probably slow -- scratch is a (channels, 256) uint16 array.
+"""
+@njit(parallel=True, fastmath=True, inline='always')
+def equalize(source, scratch, destination):
+    for i in prange(source.shape[-1]):
+        scratch[i] = np.bincount(source[..., i].flatten(), minlength=256)
+        nonzero_hist = scratch[i][scratch[i] != 0]
+        step = nonzero_hist[:-1].sum() // 255
+
+        if step == 0:
+            continue
+
+        scratch[i][1:] = scratch[i].cumsum()[:-1]
+        scratch[i] = (scratch[i] + step // 2) // step
+        scratch[i][0] = 0
+        np.clip(scratch[i], 0, 255, out=scratch[i])
+
+        # numba doesn't like 2d advanced indexing
+        for row in prange(source.shape[0]):
+            destination[row, :, i] = scratch[i][source[row, :, i]]
+
+"""
+Equalize using OpenCV -- not equivalent to
+torchvision.transforms.functional.equalize for so-far-unknown reasons.
+"""
+@njit(parallel=False, fastmath=True, inline='always')
+def fast_equalize(source, chw_scratch, destination):
+    # this seems kind of hacky
+    # also, assuming ctypes_equalize allocates a minimal amount of memory
+    # which may be incorrect -- so maybe we should do this from scratch.
+    # TODO may be a better way to do this in pure OpenCV
+    c, h, w = chw_scratch.shape
+    chw_scratch[0] = source[..., 0]
+    ctypes_equalize(chw_scratch.ctypes.data,
+                    chw_scratch.ctypes.data,
+                    h, w)
+    chw_scratch[1] = source[..., 1]
+    ctypes_equalize(chw_scratch.ctypes.data + h*w,
+                    chw_scratch.ctypes.data + h*w,
+                    h, w)
+    chw_scratch[2] = source[..., 2]
+    ctypes_equalize(chw_scratch.ctypes.data + 2*h*w,
+                    chw_scratch.ctypes.data + 2*h*w,
+                    h, w)
+    destination[..., 0] = chw_scratch[0]
+    destination[..., 1] = chw_scratch[1]
+    destination[..., 2] = chw_scratch[2]
 
 
 @njit(parallel=False, fastmath=True, inline='always')

libffcv/libffcv.cpp

@@ -74,6 +74,13 @@ extern "C" {
                    0, dest_matrix);
     }
 
+    void equalize(int64_t source_p, int64_t dest_p, int64_t sx, int64_t sy) {
+        cv::Mat source_matrix(sx, sy, CV_8U, (uint8_t*) source_p);
+        cv::Mat dest_matrix(sx, sy, CV_8U, (uint8_t*) dest_p);
+        cv::equalizeHist(source_matrix.colRange(0, sy).rowRange(0, sx),
+                        dest_matrix);
+    }
+
     void my_memcpy(void *source, void* dst, uint64_t size) {
         memcpy(dst, source, size);
     }

tests/test_rand_aug.py

@@ -13,7 +13,7 @@
 from typing import Callable, Optional, Tuple
 from ffcv.pipeline.state import State
 from ffcv.transforms.utils.fast_crop import rotate, shear, blend, \
-    adjust_contrast, posterize, invert, solarize
+    adjust_contrast, posterize, invert, solarize, equalize, fast_equalize
 import torchvision.transforms as tv
 import cv2
 import pytest
@@ -27,7 +27,7 @@ def __init__(self, size: int):
     def generate_code(self) -> Callable:
         my_range = Compiler.get_iterator()
         def randaug(im, mem):
-            dst, scratch = mem
+            dst, scratch, lut = mem
             for i in my_range(im.shape[0]):
 
                 ## TODO actual randaug logic
@@ -42,6 +42,8 @@ def randaug(im, mem):
                 ## adjust contrast
                 adjust_contrast(im[i], scratch[i][0], 0.5, dst[i])
 
+                ## equalize
+                equalize(im[i], lut[i], dst[i])
             return dst
 
         randaug.is_parallel = True
@@ -51,7 +53,8 @@ def declare_state_and_memory(self, previous_state: State) -> Tuple[State, Option
         assert previous_state.jit_mode
         return replace(previous_state, shape=(self.size, self.size, 3)), [
             AllocationQuery((self.size, self.size, 3), dtype=np.dtype('uint8')), 
-            AllocationQuery((1, self.size, self.size, 3), dtype=np.dtype('uint8'))
+            AllocationQuery((1, self.size, self.size, 3), dtype=np.dtype('uint8')),
+            AllocationQuery((3, 256), dtype=np.dtype('int16'))
         ]
 
 
@@ -191,14 +194,40 @@ def test_solarize(threshold):
     assert np.linalg.norm(Ynp.astype(np.float32) - Ych.astype(np.float32)) < 100
 
 
+def test_equalize():
+    Xnp = np.random.uniform(0, 256, size=(32, 32, 3)).astype(np.uint8)
+    #Xnp = cv2.imread('example_imgs/0249.png')
+    Xnp[5:9,5:9,:] = 0
+    Ynp = np.zeros(Xnp.shape, dtype=np.uint8)
+    #Snp_chw = np.zeros((3, 32, 32), dtype=np.uint8)
+    Snp = np.zeros((3, 256), dtype=np.int16)
+    Xch = torch.tensor(Xnp).permute(2, 0, 1)
+    Ych = tv.functional.equalize(Xch).permute(1, 2, 0).numpy()
+    #fast_equalize(Xnp, Snp_chw, Ynp)
+    equalize(Xnp, Snp, Ynp)
+
+    plt.subplot(2, 2, 1)
+    plt.imshow(Xnp)
+    plt.subplot(2, 2, 2)
+    plt.imshow(Ynp)
+    plt.subplot(2, 2, 3)
+    plt.imshow(Xch.permute(1, 2, 0).numpy())
+    plt.subplot(2, 2, 4)
+    plt.imshow(Ych)
+    plt.savefig('example_imgs/equalize.png')
+
+    assert np.linalg.norm(Ynp.astype(np.float32) - Ych.astype(np.float32)) < 100
+
+
 if __name__ == '__main__':
-    test_rotate(45)
-    test_shear(0.31)
-    test_brightness(0.5)
-    test_adjust_contrast(0.5)
-    test_posterize(2)
-    test_invert()
-    test_solarize(9)
+#     test_rotate(45)
+#     test_shear(0.31)
+#     test_brightness(0.5)
+#     test_adjust_contrast(0.5)
+#     test_posterize(2)
+#     test_invert()
+#     test_solarize(9)
+#     test_equalize()
 
     BATCH_SIZE = 512
     image_pipelines = {