-
Notifications
You must be signed in to change notification settings - Fork 7
/
feature_extract.py
158 lines (112 loc) · 6.11 KB
/
feature_extract.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
# -*- coding: utf-8 -*-
import cv2
import os, sys
import time
from scipy import misc
import argparse
import chroma
import dark_channel
import hue_disparity
import max_local_contrast
import max_local_saturation
import min_local_clr_attenuation
RESIZE_SHAPE = (128, 128, 6)
def extract_feature(args):
img_count = 0
rawimgs_src_path = args.data_path
dataset_save_path = args.save_path
rawimgs_dir = os.listdir(rawimgs_src_path)
if not os.path.exists(dataset_save_path):
os.mkdir(dataset_save_path)
for each_rawimgs_dir in rawimgs_dir:
print(each_rawimgs_dir)
each_rawimgs_path = rawimgs_src_path + '/' + each_rawimgs_dir
raw_imgs = os.listdir(each_rawimgs_path)
for each_img in raw_imgs:
raw_img_name, _ = each_img.split('.')
img_src_full_path = each_rawimgs_path + '/' + each_img
img_src_read = cv2.imread(img_src_full_path)
# resize raw imgs to 128*128
img_src_read = misc.imresize(img_src_read, [RESIZE_SHAPE[0], RESIZE_SHAPE[1]], interp='bilinear')
# ====================== Feature 1: Refined Dark Channel
if not os.path.exists(dataset_save_path + '/feature1'):
os.mkdir(dataset_save_path + '/feature1')
img_save_full_path = dataset_save_path + '/feature1/'
I = img_src_read.astype('float64') / 255.
# Patch_Size = 15*15
darkch_norm = dark_channel.DarkChannel_Norm(I, 15)
A = dark_channel.AtmLight(I, darkch_norm)
te = dark_channel.TransmissionEstimate(I, A, 15, 1)
t = dark_channel.TransmissionRefine(img_src_read, te)
darkch_refined = 1 - t
cv2.imwrite(img_save_full_path + each_img, darkch_refined * 255, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
img_count += 1
sys.stdout.write("\rProcessing a new image: %d ...... Feature maps extraction: 1 done." % img_count)
sys.stdout.flush()
# ====================== Feature 2: Max Local Contrast
if not os.path.exists(dataset_save_path + '/feature2'):
os.mkdir(dataset_save_path + '/feature2')
img_save_full_path = dataset_save_path + '/feature2/'
I = img_src_read.astype('float64') / 255.
# Patch_Size = 10*10, local_size = 5*5
contrast = max_local_contrast.Max_Local_Contrast(I)
cv2.imwrite(img_save_full_path + each_img, contrast * 255, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
sys.stdout.write("\rProcessing a new image: %d ...... Feature maps extraction: 2 done." % img_count)
sys.stdout.flush()
# ====================== Feature 3: Max Local Saturation
if not os.path.exists(dataset_save_path + '/feature3'):
os.mkdir(dataset_save_path + '/feature3')
img_save_full_path = dataset_save_path + '/feature3/'
I = img_src_read.astype('float64') / 255.
# Patch_Size = 10*10
saturation = max_local_saturation.Max_Local_Sat(I)
cv2.imwrite(img_save_full_path + each_img, saturation * 255, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
sys.stdout.write("\rProcessing a new image: %d ...... Feature maps extraction: 3 done." % img_count)
sys.stdout.flush()
# ====================== Feature 4: Min Local Color Attenuation
if not os.path.exists(dataset_save_path + '/feature4'):
os.mkdir(dataset_save_path + '/feature4')
img_save_full_path = dataset_save_path + '/feature4/'
# normalize AFTER TRANS
I = img_src_read.astype('float64')
# Patch_Size = 15*15
clr_atten = min_local_clr_attenuation.Min_Local_Clr_Atten(I)
cv2.imwrite(img_save_full_path + each_img, clr_atten * 255, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
sys.stdout.write("\rProcessing a new image: %d ...... Feature maps extraction: 4 done." % img_count)
sys.stdout.flush()
# ====================== Feature 5: Hue Disparity
if not os.path.exists(dataset_save_path + '/feature5'):
os.mkdir(dataset_save_path + '/feature5')
img_save_full_path = dataset_save_path + '/feature5/'
I = img_src_read.astype('float64') / 255.
hue = hue_disparity.Hue_Disp(I)
cv2.imwrite(img_save_full_path + each_img, hue * 255, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
sys.stdout.write("\rProcessing a new image: %d ...... Feature maps extraction: 5 done." % img_count)
sys.stdout.flush()
# ====================== Feature 6: Chroma
if not os.path.exists(dataset_save_path + '/feature6'):
os.mkdir(dataset_save_path + '/feature6')
img_save_full_path = dataset_save_path + '/feature6/'
I = img_src_read.astype('float64') / 255.
ch = chroma.Chroma(I)
cv2.imwrite(img_save_full_path + each_img, ch * 255, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
sys.stdout.write("\rProcessing a new image: %d ...... Feature maps extraction: 6 done." % img_count)
sys.stdout.flush()
print("Total processed images of %s: %d" % (each_rawimgs_dir, img_count))
# time.sleep(600)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--data-path', type=str, required=True,
help='directory containing raw images')
parser.add_argument('--save-path', type=str, required=True,
help='directory for saving processed feature images')
# parser.add_argument('--contrast-scale', type=int, default=10,
# help='scale factor for maximum local contrast feature')
# parser.add_argument('--sat-scale', type=int, default=10,
# help='scale factor for maximum local saturation feature')
# parser.add_argument('--clratten-scale', type=int, default=15,
# help='scale factor for minimum local color attenuation feature')
args = parser.parse_args()
extract_feature(args)
if __name__ == '__main__':
main()