-
Notifications
You must be signed in to change notification settings - Fork 0
/
multiview.py
119 lines (85 loc) · 4.71 KB
/
multiview.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
'''
multiview.py
create multi-view images of original images
use for image augmentation
'''
import argparse
import cv2
import numpy as np
from skimage import morphology
import os
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, default='./data/Original_img',
help='input images path.')
parser.add_argument('--output', type=str, default='./data',
help='output muti-view images path.')
parser.add_argument('--gray', action="store_true", default=False,
help='save gray_img flag: default False')
parser.add_argument('--binary', action="store_true", default=False,
help='binary flag: default False')
parser.add_argument('--blocksize', type=int, default=41,
help='blocksize for binary')
parser.add_argument('--C', type=float, default=5,
help='C for binary')
parser.add_argument('--skeletonize', action="store_true", default=False,
help='skeletonize model flag: default True')
args = parser.parse_args()
return args
def find_file(path):
file_list=os.listdir(path)
file_name_list = []
for file_name in file_list:
file_name_list.append(file_name)
print(f'find {len(file_name_list)} files (out of {len(file_list)} all files) under {path}')
return file_name_list
def find_img_file(path):
file_list=os.listdir(path)
img_name_list = []
for file_name in file_list:
if file_name[-4:]=='.png':
img_name_list.append(file_name)
print(f'find {len(img_name_list)} image files (out of {len(file_list)} all files) under {path}')
return img_name_list
def multiview(args):
img_files_list = find_file(args.input)
os.makedirs(args.output, exist_ok=True)
for img_files_name in img_files_list:
img_classification_list = find_file(os.path.join(args.input, img_files_name))
for img_classification in img_classification_list:
img_name_list = find_img_file(os.path.join(args.input, img_files_name, img_classification))
for img_name in img_name_list:
img_origin_dir = os.path.join(args.input, img_files_name, img_classification, img_name)
print(img_origin_dir)
# ---- STEP 1 ---- read origin image and convert to gray img
img_origin = cv2.imread(img_origin_dir)
img_gray = cv2.cvtColor(img_origin, cv2.COLOR_BGR2GRAY)
if args.gray:
os.makedirs(os.path.join(args.output, 'gray', img_files_name, img_classification), exist_ok=True)
gray_dir = os.path.join(args.output, 'gray', img_files_name, img_classification)
cv2.imwrite(os.path.join(gray_dir, img_name), img_gray)
# ---- STEP 2 ---- convert to binary for skeletonize-process
if args.binary:
# binary (black background)
blocksize = args.blocksize
C = args.C
img_binary = cv2.adaptiveThreshold(img_gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, blocksize, C)
# ---- STEP 3 ---- binary convert to skeleton
if args.skeletonize:
os.makedirs(os.path.join(args.output, 'skeleton', img_files_name, img_classification), exist_ok=True)
Skeleton_dir = os.path.join(args.output, 'skeleton', img_files_name, img_classification)
img_skeletoniz = img_binary
img_skeletoniz[img_skeletoniz==255] = 1
skeleton0 = morphology.skeletonize(img_skeletoniz)
skeleton = skeleton0.astype(np.uint8)*255
# convert to white background and save skeleton image
skeleton_INV = skeleton
where_0 = np.where(skeleton_INV == 0)
where_255 = np.where(skeleton_INV == 255)
skeleton_INV[where_0] = 255
skeleton_INV[where_255] = 0
cv2.imwrite(os.path.join(Skeleton_dir, img_name), skeleton_INV)
if __name__ == "__main__":
args = get_args()
print(args)
multiview(args)