-
-
Save HViktorTsoi/8e8b0468a9fb07842669aa368382a7df to your computer and use it in GitHub Desktop.
import numpy as np | |
import cv2 | |
def correction( | |
img, | |
shadow_amount_percent, shadow_tone_percent, shadow_radius, | |
highlight_amount_percent, highlight_tone_percent, highlight_radius, | |
color_percent | |
): | |
""" | |
Image Shadow / Highlight Correction. The same function as it in Photoshop / GIMP | |
:param img: input RGB image numpy array of shape (height, width, 3) | |
:param shadow_amount_percent [0.0 ~ 1.0]: Controls (separately for the highlight and shadow values in the image) how much of a correction to make. | |
:param shadow_tone_percent [0.0 ~ 1.0]: Controls the range of tones in the shadows or highlights that are modified. | |
:param shadow_radius [>0]: Controls the size of the local neighborhood around each pixel | |
:param highlight_amount_percent [0.0 ~ 1.0]: Controls (separately for the highlight and shadow values in the image) how much of a correction to make. | |
:param highlight_tone_percent [0.0 ~ 1.0]: Controls the range of tones in the shadows or highlights that are modified. | |
:param highlight_radius [>0]: Controls the size of the local neighborhood around each pixel | |
:param color_percent [-1.0 ~ 1.0]: | |
:return: | |
""" | |
shadow_tone = shadow_tone_percent * 255 | |
highlight_tone = 255 - highlight_tone_percent * 255 | |
shadow_gain = 1 + shadow_amount_percent * 6 | |
highlight_gain = 1 + highlight_amount_percent * 6 | |
# extract RGB channel | |
height, width = img.shape[:2] | |
img = img.astype(np.float) | |
img_R, img_G, img_B = img[..., 2].reshape(-1), img[..., 1].reshape(-1), img[..., 0].reshape(-1) | |
# The entire correction process is carried out in YUV space, | |
# adjust highlights/shadows in Y space, and adjust colors in UV space | |
# convert to Y channel (grey intensity) and UV channel (color) | |
img_Y = .3 * img_R + .59 * img_G + .11 * img_B | |
img_U = -img_R * .168736 - img_G * .331264 + img_B * .5 | |
img_V = img_R * .5 - img_G * .418688 - img_B * .081312 | |
# extract shadow / highlight | |
shadow_map = 255 - img_Y * 255 / shadow_tone | |
shadow_map[np.where(img_Y >= shadow_tone)] = 0 | |
highlight_map = 255 - (255 - img_Y) * 255 / (255 - highlight_tone) | |
highlight_map[np.where(img_Y <= highlight_tone)] = 0 | |
# // Gaussian blur on tone map, for smoother transition | |
if shadow_amount_percent * shadow_radius > 0: | |
# shadow_map = cv2.GaussianBlur(shadow_map.reshape(height, width), ksize=(shadow_radius, shadow_radius), sigmaX=0).reshape(-1) | |
shadow_map = cv2.blur(shadow_map.reshape(height, width), ksize=(shadow_radius, shadow_radius)).reshape(-1) | |
if highlight_amount_percent * highlight_radius > 0: | |
# highlight_map = cv2.GaussianBlur(highlight_map.reshape(height, width), ksize=(highlight_radius, highlight_radius), sigmaX=0).reshape(-1) | |
highlight_map = cv2.blur(highlight_map.reshape(height, width), ksize=(highlight_radius, highlight_radius)).reshape(-1) | |
# Tone LUT | |
t = np.arange(256) | |
LUT_shadow = (1 - np.power(1 - t * (1 / 255), shadow_gain)) * 255 | |
LUT_shadow = np.maximum(0, np.minimum(255, np.int_(LUT_shadow + .5))) | |
LUT_highlight = np.power(t * (1 / 255), highlight_gain) * 255 | |
LUT_highlight = np.maximum(0, np.minimum(255, np.int_(LUT_highlight + .5))) | |
# adjust tone | |
shadow_map = shadow_map * (1 / 255) | |
highlight_map = highlight_map * (1 / 255) | |
iH = (1 - shadow_map) * img_Y + shadow_map * LUT_shadow[np.int_(img_Y)] | |
iH = (1 - highlight_map) * iH + highlight_map * LUT_highlight[np.int_(iH)] | |
img_Y = iH | |
# adjust color | |
if color_percent != 0: | |
# color LUT | |
if color_percent > 0: | |
LUT = (1 - np.sqrt(np.arange(32768)) * (1 / 128)) * color_percent + 1 | |
else: | |
LUT = np.sqrt(np.arange(32768)) * (1 / 128) * color_percent + 1 | |
# adjust color saturation adaptively according to highlights/shadows | |
color_gain = LUT[np.int_(img_U ** 2 + img_V ** 2 + .5)] | |
w = 1 - np.minimum(2 - (shadow_map + highlight_map), 1) | |
img_U = w * img_U + (1 - w) * img_U * color_gain | |
img_V = w * img_V + (1 - w) * img_V * color_gain | |
# re convert to RGB channel | |
output_R = np.int_(img_Y + 1.402 * img_V + .5) | |
output_G = np.int_(img_Y - .34414 * img_U - .71414 * img_V + .5) | |
output_B = np.int_(img_Y + 1.772 * img_U + .5) | |
output = np.row_stack([output_B, output_G, output_R]).T.reshape(height, width, 3) | |
output = np.minimum(output, 255).astype(np.uint8) | |
return output |
Thanks a lot for sharing the code, it does its job perfectly. Although, there is one problem - it darkens highlights and lightens shadows, but it doesn't seems able to do the inverse - increase and darken them respectively. I'd like to know how this function can be modified for the sake of achieving that effect. And if it's already possible - how the inputs should be adjusted.
If you want to add rectangles etc onto the image afterwards you get the following error:
cv2.error: OpenCV: (-5:Bad argument) in function 'rectangle'
Making a copy of the image upon return on line 90 will fix this:
output = np.minimum(output, 255).astype(np.uint8).copy()
or if you changed it as per @gapeot suggestion:
output = np.minimum(np.maximum(output, 0), 255).astype(np.uint8).copy()
I think
w = 1 - np.minimum(2 - (shadow_map + highlight_map), 1)
should be changed to "w = np.minimum(2 - (shadow_map + highlight_map), 1)".