88 lines
2.8 KiB
Python
88 lines
2.8 KiB
Python
import os
|
||
import xml.etree.ElementTree as ET
|
||
import cv2
|
||
|
||
dirpath = r'E:\Train\EM14XML' # 原来存放xml文件的目录
|
||
newdir = r'E:\Train\EM14TXT' # 修改label后形成的txt目录
|
||
|
||
imagepath = r"E:\Train\EM14IMG"
|
||
|
||
if not os.path.exists(newdir):
|
||
os.makedirs(newdir)
|
||
|
||
for fp in os.listdir(dirpath):
|
||
print(fp)
|
||
print(fp.split('.')[0])
|
||
|
||
imagename = fp.split('.')[0] + '.tiff'
|
||
imagename_path = os.path.join(imagepath,imagename)
|
||
print(imagename_path)
|
||
image = cv2.imread(imagename_path)
|
||
# cv2.imshow('image',image)
|
||
# cv2.waitKey(0)
|
||
|
||
if fp == '.DS_Store':
|
||
continue
|
||
root = ET.parse(os.path.join(dirpath, fp)).getroot() #获取xml的根节点
|
||
xmin, ymin, xmax, ymax = 0, 0, 0, 0
|
||
sz = root.find('size')
|
||
print(sz)
|
||
width = float(sz[0].text)
|
||
height = float(sz[1].text)
|
||
filename = root.find('filename').text
|
||
for child in root.findall('object'): # 找到图片中的所有框
|
||
print(child)
|
||
sub = child.find('bndbox') # 找到框的标注值并进行读取
|
||
xmin = float(sub[0].text)
|
||
ymin = float(sub[1].text)
|
||
xmax = float(sub[2].text)
|
||
ymax = float(sub[3].text)
|
||
try: # 转换成yolov3的标签格式，需要归一化到（0-1）的范围内
|
||
x_center = (xmin + xmax) / (2 * width)
|
||
y_center = (ymin + ymax) / (2 * height)
|
||
w = (xmax - xmin) / width
|
||
h = (ymax - ymin) / height
|
||
except ZeroDivisionError:
|
||
print(filename, '的 width有问题')
|
||
sub_class = str(child.find('name').text)
|
||
if sub_class == 'EM17':
|
||
cls = 2
|
||
elif sub_class == 'EM14':
|
||
cls = 0
|
||
elif sub_class == 'EM170':
|
||
cls = 3
|
||
elif sub_class == 'EM18':
|
||
cls = 1
|
||
elif sub_class == 'EM19':
|
||
cls = 4
|
||
elif sub_class == 'EM190':
|
||
cls = 5
|
||
elif sub_class == 'EM20':
|
||
cls = 6
|
||
elif sub_class == 'EM200':
|
||
cls = 7
|
||
elif sub_class == 'EM201':
|
||
cls = 8
|
||
elif sub_class == 'EM202':
|
||
cls = 9
|
||
elif sub_class == 'EM203':
|
||
cls = 10
|
||
elif sub_class == 'EM180':
|
||
cls = 11
|
||
elif sub_class == 'EM141':
|
||
cls = 12
|
||
else:
|
||
continue
|
||
|
||
# cv.namedWindow('rect', 1)
|
||
# cv.imshow('rect', image)
|
||
# cv.waitKey(0)
|
||
|
||
with open(os.path.join(newdir, fp.split('.')[0] + '.txt'), 'a+') as f:
|
||
f.write(' '.join([str(cls), str(x_center), str(y_center), str(w), str(h) + '\n']))
|
||
|
||
|
||
cv2.rectangle(image, (int(xmin), int(ymin)), (int(xmax), int(ymax)), (0, 0, 255), 2)
|
||
cv2.putText(image, str(cls), (200, 200), cv2.FONT_HERSHEY_COMPLEX, 5.0, (100, 200, 200), 2)
|
||
cv2.imwrite('E:\\Train\\EM14Result\\' + imagename, image)
|