-
Notifications
You must be signed in to change notification settings - Fork 0
/
Correction.py
110 lines (91 loc) · 4.62 KB
/
Correction.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
import numpy as np
import cv2
from skimage import exposure
import os
output_folder='' # output folder
input_folder='' # input folder
def correction(
img,
shadow_amount_percent, shadow_tone_percent, shadow_radius,
highlight_amount_percent, highlight_tone_percent, highlight_radius,
color_percent,out
):
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)
logarithmic_corrected = exposure.adjust_log(output, 1)
cv2.imwrite(out,logarithmic_corrected)
return output
for root,dirs,files in os.walk(input_folder):
for file in files:
if file.endswith('.jpg') and not file.startswith(('.','_')):
filename = os.path.join(root,file)
output_filename=os.path.join(output_folder,os.path.basename(file))
opencv_image_with_bgr_channels = cv2.imread(filename,3)
# img = opencv_image_with_bgr_channels[:,:, ::-1]
img = opencv_image_with_bgr_channels
shadow_amount_percent=0.1 #0.3
shadow_tone_percent=1.0 #1
shadow_radius=1 #1
highlight_amount_percent=0 #0
highlight_tone_percent=0.1 #0.1
highlight_radius=0 #0
color_percent=0 #0
correction(
img,
shadow_amount_percent, shadow_tone_percent, shadow_radius,
highlight_amount_percent, highlight_tone_percent, highlight_radius,color_percent,output_filename
)