basic-image-editing / color_conversions.py
balthou's picture
discard division warning in HSV transform
56a9526
import numpy as np
import warnings
def rgb_to_hsv(rgb):
with warnings.catch_warnings():
# Suppress the specific RuntimeWarning
warnings.filterwarnings(
"ignore", category=RuntimeWarning, message="invalid value encountered in divide")
# Ensure the input is a float numpy array and in the range [0, 1]
rgb = np.asarray(rgb)
# Separate the R, G, B channels
r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
# Max and min of RGB values
max_rgb = np.max(rgb, axis=-1)
min_rgb = np.min(rgb, axis=-1)
delta = max_rgb - min_rgb
# Hue calculation
hue = np.zeros_like(max_rgb)
# Avoid division by zero: Only calculate where delta is non-zero
mask = delta != 0
# Red is max
idx = (max_rgb == r) & mask
hue[idx] = (60 * ((g - b) / delta % 6))[idx]
# Green is max
idx = (max_rgb == g) & mask
hue[idx] = (60 * ((b - r) / delta + 2))[idx]
# Blue is max
idx = (max_rgb == b) & mask
hue[idx] = (60 * ((r - g) / delta + 4))[idx]
# Saturation calculation
saturation = np.zeros_like(max_rgb)
saturation[mask] = delta[mask] / max_rgb[mask]
# Handle the edge case where max_rgb is 0 (black pixels)
saturation[max_rgb == 0] = 0
# Value calculation
value = max_rgb
# Stack the HSV components together
hsv = np.stack((hue, saturation, value), axis=-1)
return hsv
def hsv_to_rgb(hsv):
# Ensure the input is a float numpy array and in the range [0, 1]
hsv = np.asarray(hsv)
# Separate the H, S, V channels
h, s, v = hsv[..., 0], hsv[..., 1], hsv[..., 2]
# Chromaticity component
c = v * s
x = c * (1 - np.abs((h / 60) % 2 - 1))
m = v - c
# Initialize RGB
r, g, b = np.zeros_like(h), np.zeros_like(h), np.zeros_like(h)
# Compute the RGB components based on the hue range
h0_60 = (h >= 0) & (h < 60)
r[h0_60], g[h0_60], b[h0_60] = c[h0_60], x[h0_60], 0
h60_120 = (h >= 60) & (h < 120)
r[h60_120], g[h60_120], b[h60_120] = x[h60_120], c[h60_120], 0
h120_180 = (h >= 120) & (h < 180)
r[h120_180], g[h120_180], b[h120_180] = 0, c[h120_180], x[h120_180]
h180_240 = (h >= 180) & (h < 240)
r[h180_240], g[h180_240], b[h180_240] = 0, x[h180_240], c[h180_240]
h240_300 = (h >= 240) & (h < 300)
r[h240_300], g[h240_300], b[h240_300] = x[h240_300], 0, c[h240_300]
h300_360 = (h >= 300) & (h < 360)
r[h300_360], g[h300_360], b[h300_360] = c[h300_360], 0, x[h300_360]
# Add m to each component and stack them together
rgb = np.stack((r + m, g + m, b + m), axis=-1)
return rgb
def test_conversion():
# RGB array example (range [0, 1])
rgb = np.array([[[0.5, 0.4, 0.7], [0.1, 0.2, 0.3]],
[[0.8, 0.7, 0.6], [0.9, 0.1, 0.2]]])
# rgb = np.random.rand(1000, 100, 3)
# Convert RGB to HSV
hsv = rgb_to_hsv(rgb)
print(hsv.shape)
# Convert HSV back to RGB
rgb_back = hsv_to_rgb(hsv)
assert np.allclose(rgb, rgb_back), "Conversion failed!"
if __name__ == "__main__":
test_conversion()