目标检测图片和标签自动化处理

# -*- coding:utf-8 -*- import os import shutil import xml.etree.cElementTree as ET import tqdm from xml.dom.minidom import Document import os import cv2 from tqdm import tqdm # 将txt文件转换为对应的xml文件 def makexml(picPath, txtPath, xmlPath): if os.path.exists(xmlPath): shutil.rmtree(xmlPath) os.makedirs(xmlPath) classes = {'0': "sawclothes", '1': "other_clothes"} files = os.listdir(picPath) for i, name in enumerate(tqdm(files)): img = cv2.imread(os.path.join(picPath, name)) Pheight, Pwidth, Pdepth = img.shape # 创建XML文档对象 xmlBuilder = Document() annotation = xmlBuilder.createElement("annotation") # 创建Annotation标签 xmlBuilder.appendChild(annotation) # 添加folder标签 folder = xmlBuilder.createElement("folder") foldercontent = xmlBuilder.createTextNode("VOC") folder.appendChild(foldercontent) annotation.appendChild(folder) # 添加filename标签 filename = xmlBuilder.createElement("filename") filenamecontent = xmlBuilder.createTextNode(name) filename.appendChild(filenamecontent) annotation.appendChild(filename) # 添加size标签 size = xmlBuilder.createElement("size") width = xmlBuilder.createElement("width") widthcontent = xmlBuilder.createTextNode(str(Pwidth)) width.appendChild(widthcontent) size.appendChild(width) height = xmlBuilder.createElement("height") heightcontent = xmlBuilder.createTextNode(str(Pheight)) height.appendChild(heightcontent) size.appendChild(height) depth = xmlBuilder.createElement("depth") depthcontent = xmlBuilder.createTextNode(str(Pdepth)) depth.appendChild(depthcontent) size.appendChild(depth) annotation.appendChild(size) txt = os.path.join(txtPath, name.replace('jpg', 'txt')) if not os.path.exists(txt): # 若文本文件不存在，则直接保存XML文件并处理下一张图片 print(txt) f = open(os.path.join(xmlPath, name[0:-4] + '.xml'), 'w', encoding="utf-8") xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8') f.close() continue # 若文本文件存在，则读取文本内容并添加相应的标签 txtFile = open(txt) txtList = txtFile.readlines() for j in txtList: oneline = j.strip().split(" ") object = xmlBuilder.createElement("object") picname = xmlBuilder.createElement("name") nameContent = xmlBuilder.createTextNode(classes[oneline[0]]) picname.appendChild(nameContent) object.appendChild(picname) # 添加pose、truncated、difficult等标签 # 添加bndbox标签，并计算bounding box的坐标 # 将以上标签添加至annotation标签下 pose = xmlBuilder.createElement("pose") posecontent = xmlBuilder.createTextNode("Unspecified") pose.appendChild(posecontent) object.appendChild(pose) truncated = xmlBuilder.createElement("truncated") truncatedContent = xmlBuilder.createTextNode("0") truncated.appendChild(truncatedContent) object.appendChild(truncated) difficult = xmlBuilder.createElement("difficult") difficultContent = xmlBuilder.createTextNode("0") difficult.appendChild(difficultContent) object.appendChild(difficult) bndbox = xmlBuilder.createElement("bndbox") xmin = xmlBuilder.createElement("xmin") mathData = int(((float(oneline[1])) * Pwidth + 1) - (float(oneline[3])) * 0.5 * Pwidth) xminContent = xmlBuilder.createTextNode(str(mathData)) xmin.appendChild(xminContent) bndbox.appendChild(xmin) ymin = xmlBuilder.createElement("ymin") mathData = int(((float(oneline[2])) * Pheight + 1) - (float(oneline[4])) * 0.5 * Pheight) yminContent = xmlBuilder.createTextNode(str(mathData)) ymin.appendChild(yminContent) bndbox.appendChild(ymin) xmax = xmlBuilder.createElement("xmax") mathData = int(((float(oneline[1])) * Pwidth + 1) + (float(oneline[3])) * 0.5 * Pwidth) xmaxContent = xmlBuilder.createTextNode(str(mathData)) xmax.appendChild(xmaxContent) bndbox.appendChild(xmax) ymax = xmlBuilder.createElement("ymax") mathData = int(((float(oneline[2])) * Pheight + 1) + (float(oneline[4])) * 0.5 * Pheight) ymaxContent = xmlBuilder.createTextNode(str(mathData)) ymax.appendChild(ymaxContent) bndbox.appendChild(ymax) object.appendChild(bndbox) annotation.appendChild(object) f = open(os.path.join(xmlPath, name.replace('jpg', 'xml')), 'w') xmlBuilder.writexml(f, indent='\t', newl='\n', addindent='\t', encoding='utf-8') f.close() # 将文件夹中的图片和标注文件分开存放 def copy(filePath, imgPath, labelPath): for fileName in os.listdir(filePath): srcName = os.path.join(filePath, fileName) dstName1 = os.path.join(imgPath, fileName) dstName2 = os.path.join(labelPath, fileName) if fileName.split('.')[-1] == 'jpg': pass # shutil.copyfile(srcName, dstName1) elif fileName.split('.')[-1] == 'xml': shutil.copyfile(srcName, dstName2) # 修改img和标签名，同时对两者进行重命名 def find_imgxml_rename(img_path, label_path): # 根据图像编号修改i数据 i = 0 for img in os.listdir(img_path): # 获得图片文件的后缀名 i = i + 1 filename, extension = os.path.splitext(img) os.rename(img_path + '/' + img, img_path + '/' + 'crop_othes' + str(i) + filename + '.jpg') os.rename(label_path + '/' + filename + '.xml', label_path + '/' + 'crop_othes' + str(i) + filename + '.xml') # 删除标签 def remove_img(imgPath, labelPath): for label in os.listdir(labelPath): # print("parsing image:", os.path.join(imgPath, image)) base_name = label.split('.xml')[0] if (base_name + '.jpg') in os.listdir(imgPath): continue else: os.remove(os.path.join(labelPath, label)) # 删除小像素目标 def delete_pix(labelPath): for file in os.listdir(labelPath): if file.split('.')[-1] == 'xml': filePath = os.path.join(labelPath, file) tree = ET.parse(filePath) root = tree.getroot() for child in root.findall('object'): for name,bndbox in zip(child.findall("name"), child.findall("bndbox")): for xmin, ymin, xmax, ymax in zip(bndbox.findall("xmin"), bndbox.findall("ymin"), bndbox.findall("xmax"), bndbox.findall("ymax")): width = float(xmax.text) - float(xmin.text) height = float(ymax.text) - float(ymin.text) print('width: ', width) print('height: ', height) if width < 48 and height < 48: print('remove' + str(filePath)) root.remove(child) tree.write(filePath) # 删除不含指定标签的xml文件 def delete_labelxml(labelPath, label): for file in os.listdir(labelPath): if file.split('.')[-1] == 'xml': filePath = os.path.join(labelPath, file) tree = ET.parse(filePath) root = tree.getroot() if root.find('object') is None: print('remove' + str(filePath)) os.remove(filePath) else: for child in root.findal