forked from MasazI/crfasrnn-training
/
filter_images.py
103 lines (82 loc) · 2.6 KB
/
filter_images.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
#!/usr/bin/env python
# Martin Kersner, m.kersner@gmail.com
# 2016/01/18
from __future__ import print_function
import os
import sys
from skimage.io import imread
import numpy as np
from utils import get_id_classes, convert_from_color_segmentation
# segmentation labels file extention
ext = '.png'
# recognize cateogry
class_names = ['car', 'motorbike', 'bus']
def main():
##
#ext = '.png'
## set specific classes
#class_names = ['bird', 'bottle', 'chair']
##
path, txt_file = process_arguments(sys.argv)
# remove old files
clear_class_logs(class_names)
# get interested classes indexes
class_ids = get_id_classes(class_names)
# get from labels image list
with open(txt_file, 'rb') as f:
i = 0
for img_name in f:
# delete white space in prefix and suffix
img_name = img_name.strip()
detected_class = contain_class(os.path.join(path, img_name)+ext, class_ids, class_names)
if detected_class:
log_class(img_name, detected_class)
print("No.%d: %s --> detect class: %s" % (i, img_name, detected_class))
else:
print("No.%d: %s --> no class" % (i, img_name))
i += 1
def clear_class_logs(class_names):
for c in class_names:
file_name = c + '.txt'
if os.path.isfile(file_name):
os.remove(file_name)
def log_class(img_name, detected_class):
'''
output file each class
'''
with open(detected_class + '.txt', 'ab') as f:
print(img_name, file=f)
def contain_class(img_name, class_ids, class_names):
'''
arguments:
img_name: name of image
class_ids: interested classes index
class_names: interested classes name
'''
img = imread(img_name)
# If label is three-dimensional image we have to convert it to
# corresponding labels (0 - 20). Currently anticipated labels are from
# VOC pascal datasets.
# if img is rgb structure, transform img grayscale
if (len(img.shape) > 2):
img = convert_from_color_segmentation(img)
for i,j in enumerate(class_ids):
if j in np.unique(img):
# if image pixel have class_id, return calss name.
return class_names[i]
return False
def process_arguments(argv):
if len(argv) != 3:
help()
dataset_segmentation_path = argv[1]
list_of_images = argv[2]
return dataset_segmentation_path, list_of_images
def help():
print('Usage: python filter_images.py PATH LIST_FILE\n'
'PATH points to directory with segmentation image labels.\n'
'LIST_FILE denotes text file containing names of images in PATH.\n'
'Names do not include extension of images.'
, file=sys.stderr)
exit()
if __name__ == '__main__':
main()