-
Notifications
You must be signed in to change notification settings - Fork 0
/
convert.py
105 lines (85 loc) · 2.92 KB
/
convert.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
# -*- coding: utf-8 -*-
"""
This script is to convert the txt annotation files to appropriate format needed by YOLO
"""
import os
from os import walk, getcwd
from PIL import Image
classes = ['10','20','50','100','200','500','2000']
def convert(size, box):
dw = 1./size[0]
dh = 1./size[1]
x = (box[0] + box[1])/2.0
y = (box[2] + box[3])/2.0
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
"""-------------------------------------------------------------------"""
""" Configure Paths"""
mypath = "Labels/2000/"
outpath = "Labels/output_2000/" # where to store the final output txt file
# on this object we want to train. Remember always train one class at a time. So I have trained first on it and now I am commenting this to work on other class
cls = "2000" # on this object we want to train
if cls not in classes:
exit(0)
cls_id = classes.index(cls)
wd = getcwd()
list_file = open('%s/%s_list.txt'%(wd, cls), 'w')
""" Get input text file list """
txt_name_list = []
for (dirpath, dirnames, filenames) in walk(mypath):
txt_name_list.extend(filenames)
break
print(txt_name_list)
""" Process """
for txt_name in txt_name_list:
""" Open input text files """
txt_path = mypath + txt_name
print("Input:" + txt_path)
txt_file = open(txt_path, "r")
lines = txt_file.read().split('\n') #for ubuntu, use "\r\n" instead of "\n"
""" Open output text files """
txt_outpath = outpath + txt_name
print("Output:" + txt_outpath)
txt_outfile = open(txt_outpath, "w")
""" Convert the data to YOLO format """
ct = 0
for line in lines:
#print('lenth of line is: ')
#print(len(line))
#print('\n')
if(len(line) >= 2):
ct = ct + 1
#print(line + "\n")
print(line)
elems = line.split(' ')
print("first")
print(elems)
print("finshed first")
xmin = elems[0]
xmax = elems[2]
ymin = elems[1]
ymax = elems[3]
#
img_path = str('%s/Images/%s/%s.png'%(wd, cls, os.path.splitext(txt_name)[0]))
#t = magic.from_file(img_path)
#wh= re.search('(\d+) x (\d+)', t).groups()
im=Image.open(img_path)
w= int(im.size[0])
h= int(im.size[1])
#w = int(xmax) - int(xmin)
#h = int(ymax) - int(ymin)
# print(xmin)
print(w, h)
b = (float(xmin), float(xmax), float(ymin), float(ymax))
bb = convert((w,h), b)
print(bb)
txt_outfile.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
""" Save those images with bb into list"""
if(ct != 0):
list_file.write('%s/Images/%s/%s.JPG\n'%(wd, cls, os.path.splitext(txt_name)[0]))
list_file.close()