Faster HSV augmentation

[developer0hye]

remove datatype conversion process that can be skipped.

import random
import numpy as np
import cv2
import time

def anchor_augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    dtype = img.dtype  # uint8

    x = np.arange(0, 256, dtype=np.int16)
    lut_hue = ((x * r[0]) % 180).astype(dtype)
    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed

def proposed_augment_hsv(img, hgain=0.5, sgain=0.5, vgain=0.5):
    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    dtype = img.dtype  # uint8

    x = np.arange(0, 256, dtype=r.dtype)
    lut_hue = ((x * r[0]) % 180).astype(dtype)
    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed

img = np.random.randint(low=0, high=255, size=(640, 640, 3), dtype=np.uint8)

%timeit anchor_augment_hsv(img)
%timeit proposed_augment_hsv(img)

I measured the execution time of the two implementations on google colab.

relative execution time: anchor/proposed = 1.06

🛠️ PR Summary

_{Made with ❤️ by Ultralytics Actions}

🌟 Summary

Refinement of data type handling in image augmentation function.

📊 Key Changes

• Changed the data type initialization in the augment_hsv function from a fixed np.int16 to the dynamic data type of the input data r.dtype.
• Removed unnecessary type conversion after merging augmented channels.

🎯 Purpose & Impact

• The purpose of this change is to maintain consistency in data types and optimize the performance of the image augmentation process.
• This impacts users by potentially preventing data type mismatch errors and enhancing the efficiency of image preprocessing in machine learning workflows. 🚀
• Users with varying data types in their images will benefit from improved handling, leading to more robust and error-resistant data augmentation. 🎨

[glenn-jocher]

@developer0hye hmm interesting. I get a faster speed also:

1.47 ms ± 136 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
1.17 ms ± 34.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

I think with this augmentation we have to be very careful though, as overflows (when the values are less than 0 or more than 255) are handled differently between uint8 and floats. I'll take a look.

[glenn-jocher]

@developer0hye we probably want a new function that uses the same random values r both ways and compares them using np.allclose(). If we run it a few hundred times and the results are always allclose True then we can be sure we aren't modifying the functionality.

[glenn-jocher]

@developer0hye allclose seems True all the time!

import copy

import cv2
import numpy as np


def anchor_augment_hsv(img, r, hgain=0.5, sgain=0.5, vgain=0.5):
    # r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    dtype = img.dtype  # uint8

    x = np.arange(0, 256, dtype=np.int16)
    lut_hue = ((x * r[0]) % 180).astype(dtype)
    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val))).astype(dtype)
    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed


def proposed_augment_hsv(img, r, hgain=0.5, sgain=0.5, vgain=0.5):
    # r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains
    hue, sat, val = cv2.split(cv2.cvtColor(img, cv2.COLOR_BGR2HSV))
    dtype = img.dtype  # uint8

    x = np.arange(0, 256, dtype=r.dtype)
    lut_hue = ((x * r[0]) % 180).astype(dtype)
    lut_sat = np.clip(x * r[1], 0, 255).astype(dtype)
    lut_val = np.clip(x * r[2], 0, 255).astype(dtype)

    img_hsv = cv2.merge((cv2.LUT(hue, lut_hue), cv2.LUT(sat, lut_sat), cv2.LUT(val, lut_val)))
    cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)  # no return needed


hgain, sgain, vgain = 0.5, 0.5, 0.5
for _ in range(100):
    img = np.random.randint(low=0, high=255, size=(640, 640, 3), dtype=np.uint8)
    im2 = copy.deepcopy(img)
    r = np.random.uniform(-1, 1, 3) * [hgain, sgain, vgain] + 1  # random gains

    anchor_augment_hsv(img, r)
    proposed_augment_hsv(im2, r)

    print(np.allclose(img, im2))

[glenn-jocher]

@developer0hye PR is merged. Thank you for your contributions!

A millisecond may not seem like much, but when the function is used thousands of times the speed improvement can be very significant!

[developer0hye]

@glenn-jocher
Yeah! That's one small step for faster hsv augmentation, one giant leap for all epochs with large batchsize.