Last active
June 1, 2020 13:49
-
-
Save hsuRush/f6689a62a740f0cf3ffa3ff99a79b7df to your computer and use it in GitHub Desktop.
for darknet
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import xml.etree.ElementTree as ET | |
import os | |
from os import listdir | |
from os.path import join | |
import random | |
path_to_dataset = '/data1000G/steven/lpDataSet/AOLP/all_train_sets/' | |
# images are at path_to_dataset + '/images/' | |
# the voc labels are at path_to_dataset + '/labels_voc/' | |
# the filename of images and labels should be exactly same. | |
# for example: | |
# path_to_dataset | |
# │ | |
# └───images | |
# │ │ | |
# │ │ file001.jpg | |
# │ │ file002.jpg | |
# │ └ ... | |
# │ | |
# └───labels_voc | |
# │ │ file001.xml | |
# │ │ file002.xml | |
classes = ["plate"] | |
valid_split = 0.1 | |
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) | |
def convert_annotation(image_id): | |
in_file = open(path_to_dataset + 'labels_voc/%s.xml'%(image_id)) #輸入路徑 | |
out_file = open(path_to_dataset + 'labels/%s.txt'%(image_id), 'w') #輸出路徑 | |
tree=ET.parse(in_file) #得到xml樹 | |
root = tree.getroot() #得到根 | |
size = root.find('size') #通過size標籤得到尺寸信息 | |
w = int(size.find('width').text) #分別得到照片的寬和高 | |
h = int(size.find('height').text) | |
for obj in root.iter('object'): #查找到每一個標籤對象 | |
difficult = obj.find('difficult').text #獲得difficult標籤內容 | |
cls = obj.find('name').text #獲得name標籤內容 | |
if cls not in classes or int(difficult) == 1: | |
continue | |
cls_id = classes.index(cls) | |
xmlbox = obj.find('bndbox') #獲取boundbox | |
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text)) | |
bb = convert((w,h), b) | |
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n') #給bb中數之間添加空格,並輸出 | |
if not os.path.exists(path_to_dataset + 'labels/'): #如果不存在要輸出的文件夾,則先創建 | |
os.makedirs(path_to_dataset + 'labels/') | |
image_ids = [f[:-4] for f in os.listdir(path_to_dataset + 'images/') if os.path.isfile(os.path.join(path_to_dataset + 'images/', f)) and ('jpg' in f)] | |
for image_id in image_ids: | |
convert_annotation(image_id) | |
random.shuffle(image_ids) #打亂文件名 | |
train_data_num = int((1-valid_split)*len(image_ids)) | |
train_ids = image_ids[0:train_data_num] #取前384個為訓練集 | |
train_list = open('train.txt', 'w') #打開要輸出的圖片路徑信息 | |
for train_id in train_ids: | |
train_list.write(path_to_dataset + 'images/%s.jpg\n'%(train_id)) #將訓練圖片的路徑寫入文件中 | |
train_list.close() | |
if not ((valid_split == 0.0) or (valid_split == 0)): | |
test_ids = image_ids[train_data_num:len(image_ids)] #剩餘的為測試集 | |
test_list = open('test.txt', 'w') #打開要輸出的圖片路徑信息 | |
for test_id in test_ids: | |
test_list.write(path_to_dataset + '/images/%s.jpg\n'%(test_id)) #將測試圖片的路徑寫入文件中 | |
test_list.close() | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment