Skip to content

Commit fc002f8

Browse files
coderzbxzhangbenxing
coderzbx
authored and
zhangbenxing
committed
增加cityscapes数据集的处理脚本
1 parent 2cd0169 commit fc002f8

File tree

6 files changed

+624
-0
lines changed

6 files changed

+624
-0
lines changed

‎SelfScripts/annotation_transform.py

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
# -*-coding:utf-8-*-
2+
3+
from PIL import Image
4+
import cv2
5+
import argparse
6+
import os
7+
import time
8+
9+
import multiprocessing
10+
from multiprocessing import Queue
11+
12+
from self_label import labels
13+
14+
# label_transform：分类合并的dictionary
15+
# unlabeled取dictionary最后一个key
16+
# src_dir：要转换的标注数据，dictionary的key对应的图片
17+
# dest_dir：要保存转换后的标注数据，dictionary的value对应的图片
18+
19+
class Task:
20+
def __init__(self, input_file, output_file):
21+
self.input_file = input_file
22+
self.output_file = output_file
23+
24+
25+
class LabelTransform:
26+
def __init__(self):
27+
self.label_transform = dict()
28+
29+
for l in labels:
30+
map_id = l.trainId
31+
if map_id > 18 or map_id < 0:
32+
map_id = 19
33+
self.label_transform[l.id] = map_id
34+
35+
self.queue = Queue()
36+
self.unlabeled = 19
37+
38+
print(self.label_transform)
39+
40+
return
41+
42+
def enter_queue(self, input_dir, output_dir):
43+
origin_list = os.listdir(input_dir)
44+
for _image in origin_list:
45+
image_path = os.path.join(input_dir, _image)
46+
result_path = os.path.join(output_dir, _image)
47+
name_list = _image.split('.')
48+
if len(name_list) < 2:
49+
print(image_path)
50+
continue
51+
52+
ext_name = name_list[1]
53+
if ext_name != 'png' and ext_name != 'jpg':
54+
continue
55+
56+
if os.path.exists(result_path):
57+
continue
58+
59+
task = Task(image_path, result_path)
60+
self.queue.put(task)
61+
62+
def transform(self):
63+
while not self.queue.empty():
64+
task = self.queue.get()
65+
image_path = task.input_file
66+
result_path = task.output_file
67+
68+
print(image_path)
69+
70+
img = cv2.imread(image_path)
71+
72+
width = img.shape[1]
73+
height = img.shape[0]
74+
75+
anna_img = Image.new('L', (width, height))
76+
77+
img_data = anna_img.load()
78+
for x in range(width):
79+
for y in range(height):
80+
cls_id = img[y, x][0]
81+
if cls_id in self.label_transform:
82+
label_id = self.label_transform[cls_id]
83+
img_data[x, y] = label_id
84+
else:
85+
img_data[x, y] = self.unlabeled
86+
87+
anna_img.save(result_path)
88+
89+
90+
if __name__ == '__main__':
91+
92+
parser = argparse.ArgumentParser()
93+
parser.add_argument('--src_dir', type=str, required=True)
94+
parser.add_argument('--dest_dir', type=str, required=True)
95+
parser.add_argument('--cpu', type=int)
96+
args = parser.parse_args()
97+
98+
src_dir = args.src_dir
99+
dest_dir = args.dest_dir
100+
101+
transFormer = LabelTransform()
102+
103+
if not os.path.exists(dest_dir):
104+
os.makedirs(dest_dir)
105+
106+
transFormer.enter_queue(src_dir, dest_dir)
107+
# transFormer.transform()
108+
109+
cpu_count = args.cpu
110+
if cpu_count == 0:
111+
cpu_count = 8
112+
113+
time1 = time.time()
114+
print("start processing....\n")
115+
116+
processes = []
117+
for i in range(0, cpu_count):
118+
process_ = multiprocessing.Process(target=transFormer.transform)
119+
process_.daemon = True
120+
processes.append(process_)
121+
122+
for i in range(0, cpu_count):
123+
proc_ = processes[i]
124+
if isinstance(proc_, multiprocessing.Process):
125+
proc_.start()
126+
127+
for i in range(0, cpu_count):
128+
proc_ = processes[i]
129+
if isinstance(proc_, multiprocessing.Process):
130+
proc_.join()
131+
132+
time2 = time.time()
133+
134+
print("finish in {} s\n".format(time2 - time1))

‎SelfScripts/filter_v1.txt

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
0,12,-1
2+
1,12,-1
3+
2,5,1
4+
3,8,1
5+
4,1,0
6+
5,12,-1
7+
6,1,1
8+
7,4,0
9+
8,3,0
10+
9,5,0
11+
10,5,0
12+
11,5,0
13+
12,4,0
14+
13,4,1
15+
14,4,1
16+
15,5,0
17+
16,1,1
18+
17,1,1
19+
18,1,1
20+
19,10,0
21+
20,11,0
22+
21,11,1
23+
22,12,-1
24+
23,3,0
25+
24,3,1
26+
25,1,1
27+
26,12,-1
28+
27,0,1
29+
28,12,-1
30+
29,12,-1
31+
30,6,1
32+
31,12,-1
33+
32,7,1
34+
33,5,0
35+
34,5,0
36+
35,7,1
37+
36,12,-1
38+
37,12,-1
39+
38,5,0
40+
39,12,-1
41+
40,5,0
42+
41,12,-1
43+
42,1,0
44+
43,12,-1
45+
44,2,1
46+
45,2,1
47+
46,2,1
48+
47,2,1
49+
48,2,0
50+
49,7,1
51+
50,7,1
52+
51,5,0
53+
52,11,0
54+
53,12,-1
55+
54,9,1
56+
55,9,1
57+
56,9,1
58+
57,11,1
59+
58,12,-1
60+
59,9,1
61+
60,9,1
62+
61,9,1
63+
62,9,0
64+
63,12,1
65+
64,12,1
66+
65,12,-1

‎SelfScripts/format_train.py

Lines changed: 62 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,62 @@
1+
import os
2+
import cv2
3+
import argparse
4+
5+
6+
class FormatTrainSet:
7+
def __init__(self):
8+
return
9+
10+
def format(self, image_dir, label_dir, txt_dir, image_type):
11+
image_files = os.listdir(image_dir)
12+
13+
images = []
14+
annots = []
15+
16+
for id_ in image_files:
17+
file_name = id_.split('.')
18+
file_ex = file_name[1]
19+
if file_ex != 'png' and file_ex != 'jpg':
20+
continue
21+
file_name = file_name[0]
22+
file_name_list = file_name.split("_")
23+
file_id = ''
24+
file_name_list = file_name_list[:-1]
25+
for name in file_name_list:
26+
file_id += name + "_"
27+
28+
image_name = file_id + "leftImg8bit.png"
29+
label_name = file_id + "gtFine_labelIds.png"
30+
31+
images.append(os.path.join(image_dir, image_name))
32+
annots.append(os.path.join(label_dir, label_name))
33+
34+
images.sort()
35+
annots.sort()
36+
image_count = len(images)
37+
label_count = len(annots)
38+
39+
train_txt = os.path.join(txt_dir, '{}.txt'.format(image_type))
40+
with open(train_txt, 'wb') as f:
41+
if image_count == label_count:
42+
for image, annot in zip(images, annots):
43+
str = image + ' ' + annot + '\n'
44+
f.write(str.encode("UTF-8"))
45+
46+
47+
if __name__ == '__main__':
48+
parser = argparse.ArgumentParser()
49+
parser.add_argument('--image_dir', type=str, required=False)
50+
parser.add_argument('--annot_dir', type=str, required=False)
51+
parser.add_argument('--text_dir', type=str, required=False)
52+
args = parser.parse_args()
53+
54+
55+
handle = FormatTrainSet()
56+
image_dir = args.image_dir
57+
annot_dir = args.annot_dir
58+
txt_dir = args.text_dir
59+
image_type = 'train'
60+
handle.format(image_dir=image_dir, label_dir=annot_dir, txt_dir=txt_dir, image_type=image_type)
61+
62+

0 commit comments

Comments
 (0)