-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathdehaze.py
161 lines (138 loc) · 6.52 KB
/
dehaze.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
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
# rewrite source code to python version
# -*- coding: utf-8 -*-
import numpy as np
import cv2
import argparse
def calc_y_channel(src):
tmp = src.copy()
cv2.cvtColor(tmp, cv2.COLOR_RGB2GRAY, tmp)
y_channel, u_channel, v_channel = cv2.split(tmp)
# cv2.imshow("y_channel", y_channel)
return y_channel
def calc_dark_channel(src, block_size):
b, g, r = cv2.split(src)
rgb_min = cv2.min(cv2.min(b, g), r)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (block_size, block_size))
dark_channel = cv2.erode(rgb_min, kernel)
# cv2.imshow("rgb_min", rgb_min)
# cv2.imshow("dark_channel", dark_channel)
return dark_channel
def calc_morphology_threshold_image(dark_channel, morphology_size):
# filtered the 90% pixel in dark channel
pixel_list = []
for i, j in np.ndindex(dark_channel.shape):
pixel_list.append(dark_channel[i, j])
# set the rest 10% pixel to air area
sorted_pixel_list = sorted(pixel_list, reverse=True)
threshold_sort_pixel = sorted_pixel_list[int(dark_channel.shape[0] * dark_channel.shape[1] * 0.1)]
# create a threshold image (white = air light in concern area, black = not in concern area)
threshold_img = dark_channel.copy()
for i, j in np.ndindex(dark_channel.shape):
if dark_channel[i, j] < threshold_sort_pixel:
threshold_img[i, j] = 0
else:
threshold_img[i, j] = 255
# cv2.imshow("threshold_img", threshold_img)
# morphological transformation
element = cv2.getStructuringElement(cv2.MORPH_RECT, (morphology_size, morphology_size))
morphology_threshold_img = cv2.morphologyEx(threshold_img, cv2.MORPH_OPEN, element)
return morphology_threshold_img
def calc_air_light(src, y_channel, dark_channel, threshold_img, is_show_circled_image):
circle_img = src.copy()
# find the brightest pixel in dark channel under the air area
air_light_val = 0
air_light_idx = []
for i, j in np.ndindex(dark_channel.shape):
if threshold_img[i, j] == 255:
_, air_light_val, _, air_light_idx = cv2.minMaxLoc(dark_channel)
cv2.circle(circle_img, air_light_idx, 5, (0, 255, 0), thickness=2)
# find the brightest pixel in Y channel
if is_show_circled_image is True:
_, _, _, max_pt = cv2.minMaxLoc(y_channel)
cv2.circle(circle_img, max_pt, 5, (0, 0, 255), thickness=2)
cv2.imshow("circle_point", circle_img)
return air_light_val
def calc_average_pixel(src):
min_val, max_val, _, _ = cv2.minMaxLoc(src)
average = (min_val + max_val) / 2
return average
def calc_gamma_correction(src, gamma):
# inv_gamma = 1.0 / gamma
gamma_table = [np.power(i / 255.0, gamma) * 255.0 for i in range(256)]
gamma_table = np.round(np.array(gamma_table)).astype(np.uint8)
return cv2.LUT(src, gamma_table)
def calc_transmission(min_med_img, air_light, th_omega):
m = calc_average_pixel(min_med_img) / 255.0
p = 1.3
q = 1 + (m - 0.5)
omega = min(m * p * q, th_omega)
transmission_map = np.uint8(255 * (1 - omega * min_med_img / air_light))
calc_gamma_correction(transmission_map, p - m)
# cv2.imshow("transmission_map", transmission_map)
return transmission_map
def clamp(minimum, x, maximum):
return max(minimum, min(x, maximum))
def dehaze(src, transmission, air_light):
restored_image = np.zeros(src.shape[:3], np.uint8)
for i in range(src.shape[0]):
for j in range(src.shape[1]):
t_max = max(transmission[i, j] / 255, 0.1)
for ch in range(src.shape[2]):
restored_image.itemset((i, j, ch), clamp(0, (src.item(i, j, ch) - air_light) / t_max + air_light, 255))
return restored_image
def main(input_file, file_type):
block_size = 5
morphology_transform_kernel_size = 15
median_filter_kernel_size = 5
show_circled_image = True
if file_type == 'image':
image = cv2.imread(input_file)
y_channel = calc_y_channel(image)
dark_channel = calc_dark_channel(image, block_size)
threshold_img = calc_morphology_threshold_image(dark_channel, morphology_transform_kernel_size)
air_light = calc_air_light(image, y_channel, dark_channel, threshold_img, show_circled_image)
min_med_img = cv2.medianBlur(y_channel, median_filter_kernel_size)
transmission_map = calc_transmission(min_med_img, air_light, 0.95)
dehaze_image = dehaze(image, transmission_map, air_light)
print("atmospheric light: ", air_light)
cv2.imshow("src", image)
cv2.imshow("dst", dehaze_image)
cv2.waitKey(0)
elif file_type == 'video':
cap = cv2.VideoCapture(input_file)
fps = cap.get(cv2.CAP_PROP_FPS)
frame_all = cap.get(cv2.CAP_PROP_FRAME_COUNT)
size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
video_writer = cv2.VideoWriter("output.mp4", cv2.VideoWriter_fourcc('m', 'p', '4', 'v'), fps, size)
print("[INFO] 视频FPS: {}".format(fps))
print("[INFO] 视频总帧数: {}".format(frame_all))
print("[INFO] 视频时长: {}s".format(frame_all / fps))
frame_count = 1
while frame_count < frame_all:
ret, image = cap.read()
print("画像を読み込みました！")
y_channel = calc_y_channel(image)
dark_channel = calc_dark_channel(image, block_size)
threshold_img = calc_morphology_threshold_image(dark_channel, morphology_transform_kernel_size)
air_light = calc_air_light(image, y_channel, dark_channel, threshold_img, show_circled_image)
print("Has been calculated air light")
min_med_img = cv2.medianBlur(y_channel, median_filter_kernel_size)
transmission_map = calc_transmission(min_med_img, air_light, 0.95)
dehaze_image = dehaze(image, transmission_map, air_light)
print("霧除去できました！")
# cv2.imshow("Output", image)
video_writer.write(dehaze_image)
if cv2.waitKey(25) & 0xFF == ord('q'):
cv2.destroyAllWindows()
break
frame_count += 1
print("処理数: {}".format(frame_count))
print("[INFO] 剩余数: {}".format(frame_all - frame_count))
cap.release()
cv2.destroyAllWindows()
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Haze remove for a single image.')
parser.add_argument('input', help='input file')
parser.add_argument('filetype', help='image or video')
args = parser.parse_args()
main(args.input, args.filetype)