/
enhance.py
89 lines (78 loc) · 2.81 KB
/
enhance.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
from skimage.io import imread, imsave
import numpy as np
from skimage.color import rgb2gray, rgb2yuv, yuv2rgb
from skimage import exposure, img_as_float, img_as_ubyte
from scipy.ndimage.filters import gaussian_filter
import os
import sys
import bisect
import argparse
def stretch(image, down = 0.0, up = 1.0):
img = image.copy()
img -= down
img /= (up - down)
return np.clip(img, 0, 1)
def get_down_up(img, thrs=(0.7, 1.0)):
hist = exposure.histogram(img)[0]
histmax = np.argmax(hist)
down = histmax * thrs[0]
up = histmax * thrs[1]
return down, up
def auto_white_balance_chnl(img, thrs=(0.7, 1.0)):
h, w = img.shape
down, up = get_down_up(img, thrs=thrs)
down /= 255.0
up /= 255.0
result = stretch(img, down, up)
return result
def auto_white_balance(img, thrs=(0.7, 1.0)):
result = img.copy()
if len(img.shape) == 2:
result = auto_white_balance_chnl(result, thrs=thrs)
return result
for chnl in range(img.shape[2]):
result[:, :, chnl] = auto_white_balance_chnl(result[:, :, chnl], thrs=thrs)
return result
def normalize_lighting(img, sigma=40):
blurred = gaussian_filter(img, sigma)
result = img / (blurred + 1e-7)
result -= result.min()
result /= result.max()
return np.clip(result, 0, 1)
def change_extension(filename, extension='png'):
return filename.split('.')[-2] + '.' + extension
def main():
parser = argparse.ArgumentParser(description='Image enchancement.\
Restoration of lighcting and contrast stretching.')
parser.add_argument('path', nargs='+',
help='Input filename')
parser.add_argument('-g', '--gray', action='store_true',
help='Make a grayscale image')
parser.add_argument('-s', '--sigma', type=float,
help='Value of sigma for gaussian filter')
parser.add_argument('-o', nargs='+',
help='Output filename')
parser.add_argument('-t', '--thrs', nargs=2, type=float,
default=(0.7, 1.0),
help='The low and high thresholds for contrast stretching')
parser.add_argument('-f', '--format',
help='Output image extension')
args = parser.parse_args()
outpaths = args.path
if args.o is not None and len(args.o) == len(args.path):
outpaths = args.o
if args.format is not None:
outpaths = list(map(lambda x: change_extension(x, args.format), outpaths))
sigma = args.sigma
if args.sigma is None:
sigma = imread(args.path[0]).shape[0] / 30
for img_path, out_path in zip(args.path, outpaths):
img = img_as_float(imread(img_path))
if args.gray:
img = rgb2gray(img)
#img = normalize_lighting(img, sigma)
img = auto_white_balance(img, args.thrs)
imsave(out_path, img)
print(img_path + ": done")
if __name__ == "__main__":
main()