114 lines
4.4 KiB
Python
Executable File
114 lines
4.4 KiB
Python
Executable File
'''
|
||
author ： Hatcher W
|
||
Email： hcwang@semi.ac.cn
|
||
use: 用于检测voc数据集是否正确，
|
||
input:（图像的文件夹路径，图像标签文件夹路径）
|
||
output：check路径，所有的可视化结果都在里面
|
||
'''
|
||
import os
|
||
import cv2
|
||
import xml.dom.minidom
|
||
from tqdm import tqdm
|
||
from lxml import etree
|
||
image_path="./JPEGImages/"
|
||
annotation_path="/data/Hatcher/Fine_HC/anno_xml/"
|
||
save_path ="./check/"
|
||
num = 0
|
||
files_name = os.listdir(annotation_path)
|
||
class GEN_Annotations:
|
||
def __init__(self, filename):
|
||
self.root = etree.Element("annotation")
|
||
|
||
child1 = etree.SubElement(self.root, "folder")
|
||
child1.text = "Hatcher"
|
||
|
||
child2 = etree.SubElement(self.root, "filename")
|
||
child2.text = filename
|
||
|
||
child3 = etree.SubElement(self.root, "source")
|
||
|
||
child4 = etree.SubElement(child3, "annotation")
|
||
child4.text = "MultiModel"
|
||
child5 = etree.SubElement(child3, "database")
|
||
child5.text = "Unknown"
|
||
|
||
child6 = etree.SubElement(child3, "image")
|
||
child6.text = "flickr"
|
||
child7 = etree.SubElement(child3, "flickrid")
|
||
child7.text = "35435"
|
||
|
||
def set_size(self, witdh, height, channel):
|
||
size = etree.SubElement(self.root, "size")
|
||
widthn = etree.SubElement(size, "width")
|
||
widthn.text = str(witdh)
|
||
heightn = etree.SubElement(size, "height")
|
||
heightn.text = str(height)
|
||
channeln = etree.SubElement(size, "depth")
|
||
channeln.text = str(channel)
|
||
|
||
def savefile(self, filename):
|
||
tree = etree.ElementTree(self.root)
|
||
tree.write(filename, pretty_print=True, xml_declaration=False, encoding='utf-8')
|
||
|
||
def add_pic_attr(self, label, xmin, ymin, xmax, ymax):
|
||
object = etree.SubElement(self.root, "object")
|
||
namen = etree.SubElement(object, "name")
|
||
namen.text = label
|
||
bndbox = etree.SubElement(object, "bndbox")
|
||
xminn = etree.SubElement(bndbox, "xmin")
|
||
xminn.text = str(xmin)
|
||
yminn = etree.SubElement(bndbox, "ymin")
|
||
yminn.text = str(ymin)
|
||
xmaxn = etree.SubElement(bndbox, "xmax")
|
||
xmaxn.text = str(xmax)
|
||
ymaxn = etree.SubElement(bndbox, "ymax")
|
||
ymaxn.text = str(ymax)
|
||
|
||
|
||
def make_xml(image_path,shape0,shape1,x_left, y_left, x_right, y_right):
|
||
|
||
|
||
print(image_path)
|
||
print("\n")
|
||
anno = GEN_Annotations(image_path)
|
||
anno.set_size(shape0, shape1, 3)
|
||
anno.add_pic_attr("bird", x_left, y_left, x_right, y_right)
|
||
anno.savefile(image_path)
|
||
for filename_ in tqdm(files_name):
|
||
xml_path = os.path.join(annotation_path,filename_)
|
||
try:
|
||
dom = xml.dom.minidom.parse(xml_path)
|
||
except:
|
||
#os.remove(img_path)
|
||
continue
|
||
root = dom.documentElement
|
||
objects=dom.getElementsByTagName("object")
|
||
shapes = dom.getElementsByTagName("size")
|
||
for shape in shapes:
|
||
width = shape.getElementsByTagName('width')[0].childNodes[0].data
|
||
height = shape.getElementsByTagName('height')[0].childNodes[0].data
|
||
for object in objects:
|
||
bndbox = object.getElementsByTagName('bndbox')[0]
|
||
xmin = bndbox.getElementsByTagName('xmin')[0]
|
||
ymin = bndbox.getElementsByTagName('ymin')[0]
|
||
xmax = bndbox.getElementsByTagName('xmax')[0]
|
||
ymax = bndbox.getElementsByTagName('ymax')[0]
|
||
#print(xmin.childNodes[0].data)
|
||
#print(type(xmin.childNodes[0].data))
|
||
xmin_data=int(float(xmin.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
|
||
ymin_data=int(float(ymin.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
|
||
xmax_data=int(float(xmax.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
|
||
ymax_data=int(float(ymax.childNodes[0].data.replace("tensor(","").replace(", device='cuda:0')","")))
|
||
label_name=object.getElementsByTagName('name')[0].childNodes[0].data
|
||
# cv2.rectangle(img,(xmin_data,ymin_data),(xmax_data,ymax_data),(55,255,155),1)
|
||
# num +=1
|
||
# cv2.putText(img,label_name,(xmin_data,ymin_data-3),cv2.FONT_HERSHEY_SIMPLEX,0.5,(0,255,0),2)
|
||
make_xml(xml_path, width, height, xmin_data, ymin_data, xmax_data, ymax_data)
|
||
flag=0
|
||
# flag=cv2.imwrite(save_path+"{}.jpg".format(filename),img)
|
||
# if not (flag):
|
||
# print(filename,"error")
|
||
print("all done ====================================")
|
||
print("number = ",num)
|
||
|
||
#(int((xmin_data+xmax_data/2)),int((ymin_data+ymax_data)/2)) |