diff --git a/create_labeles/add_xml.py b/create_labeles/add_xml.py new file mode 100644 index 0000000..007e5a6 --- /dev/null +++ b/create_labeles/add_xml.py @@ -0,0 +1,95 @@ + +from xml.etree.ElementTree import ElementTree, Element + + +# xml换行 + + +def indent(elem, level=0): + i = "\n" + level*"\t" + if len(elem): + if not elem.text or not elem.text.strip(): + elem.text = i + "\t" + if not elem.tail or not elem.tail.strip(): + elem.tail = i + for elem in elem: + indent(elem, level+1) + if not elem.tail or not elem.tail.strip(): + elem.tail = i + else: + if level and (not elem.tail or not elem.tail.strip()): + elem.tail = i + + +def add_xml(inforsDict,xmlFilePath): + + result = inforsDict + + for re in result: + # if re['score'] > 0.5: + + # 获得标注信息 + ObjName = list(re.keys())[0] + xmin = int(list(re.values())[0][0]) + ymin = int(list(re.values())[0][1]) + xmax = int(list(re.values())[0][2]) + ymax = int(list(re.values())[0][3]) + # xmax = xmin + r + # ymax = ymin + z + + #if ObjName == 'person': + + tree = ElementTree() + tree.parse(xmlFilePath) + + # 得到根目录 + root = tree.getroot() + + # 创建一级目录 + elementOjb = Element('object') + + elementBox = Element('bndbox') + + # 创建二级目录 + one = Element('name') + one.text = ObjName # 二级目录的值 #结果展示:1 + elementOjb.append(one) # 将二级目录加到一级目录里 + + two = Element('pose') + two.text = "Unspecified" + elementOjb.append(two) + + three = Element('truncated') + three.text = "0" + elementOjb.append(three) + + four = Element('difficult') + four.text = "0" + elementOjb.append(four) + + five = Element('xmin') + five.text = str(xmin) + elementBox.append(five) + + six = Element('xmax') + six.text = str(xmax) + elementBox.append(six) + + seven = Element('ymin') + seven.text = str(ymin) + elementBox.append(seven) + + eight = Element('ymax') + eight.text = str(ymax) + elementBox.append(eight) + + # 将一级目录加到根目录里 + elementOjb.append(elementBox) + root.append(elementOjb) + # 换行缩进 + indent(elementOjb) + indent(elementBox) + # 让结果保存进文件就可以了 + tree.write(xmlFilePath, encoding='utf-8', xml_declaration=True) + + diff --git a/create_labeles/create_xml.py b/create_labeles/create_xml.py new file mode 100644 index 0000000..1ce79c2 --- /dev/null +++ b/create_labeles/create_xml.py @@ -0,0 +1,58 @@ + +from lxml.etree import Element, SubElement, tostring + + +def create_xml(boxs,img_shape,xml_path): + """ + 创建xml文件,依次写入xml文件必备关键字 + :param boxs: txt文件中的box + :param img_shape: 图片信息,xml中需要写入WHC + :return: + """ + node_root = Element('annotation') + node_folder = SubElement(node_root, 'folder') + node_folder.text = 'Images' + node_filename = SubElement(node_root, 'filename') + node_filename.text = str(img_shape[3]) + node_size = SubElement(node_root, 'size') + node_width = SubElement(node_size, 'width') + node_width.text = str(img_shape[1]) + node_height = SubElement(node_size, 'height') + node_height.text = str(img_shape[0]) + node_depth = SubElement(node_size, 'depth') + node_depth.text = str(img_shape[2]) + + if len(boxs)>=1: # 循环写入box + for box in boxs: + node_object = SubElement(node_root, 'object') + node_name = SubElement(node_object, 'name') + # if str(list_[4]) == "person": # 根据条件筛选需要标注的标签,例如这里只标记person这类,不符合则直接跳过 + # node_name.text = str(list_[4]) + # else: + # continue + node_name.text = str(list(box.keys())[0]) + node_difficult = SubElement(node_object, 'difficult') + node_difficult.text = '0' + node_bndbox = SubElement(node_object, 'bndbox') + node_xmin = SubElement(node_bndbox, 'xmin') + node_xmin.text = str(int(list(box.values())[0][0])) + node_ymin = SubElement(node_bndbox, 'ymin') + node_ymin.text = str(int(list(box.values())[0][1])) + node_xmax = SubElement(node_bndbox, 'xmax') + node_xmax.text = str(int(list(box.values())[0][2])) + node_ymax = SubElement(node_bndbox, 'ymax') + node_ymax.text = str(int(list(box.values())[0][3])) + + xml = tostring(node_root, pretty_print=True) # 格式化显示,该换行的换行 + + file_name = img_shape[3].split(".")[0] + filename = xml_path+"/{}.xml".format(file_name) + + f = open(filename, "wb") + f.write(xml) + f.close() + + + + + diff --git a/create_labeles/model_to_xml.py b/create_labeles/model_to_xml.py new file mode 100644 index 0000000..5bf57f9 --- /dev/null +++ b/create_labeles/model_to_xml.py @@ -0,0 +1,44 @@ +from add_xml import add_xml +from create_xml import create_xml +from xz_yolov8_atm.yolov8_det import analysis_yolov8 +import glob +import cv2 +import os +from ultralytics import YOLO + + +if __name__ == '__main__': + + imgpath = 'E:/BANK_xz_all/imgpath' + xmlpath = 'E:/BANK_xz_all/xz_yolov8_atm/output' + + imglist = os.listdir(imgpath) + xmllist = os.listdir(xmlpath) + + model_all = YOLO('xz_yolov8_atm/model_files/best(3)_all.pt') + + for img in imglist: + + images_path = os.path.join(imgpath,img) + print(images_path) + xml_path = os.path.join(xmlpath,img.split('.')[0] + '.xml') + print(xml_path) + + images = cv2.imread(images_path) + # print(images) + + results = analysis_yolov8(images=images, + model_coco=model_all, + confidence=0.5) + + w,h,d = images.shape + img_shape = (w,h,d,img) + + if os.path.isfile(xml_path): + + add_labels = add_xml(inforsDict=results, + xmlFilePath=xml_path) + else: + create_new = create_xml(boxs=results, + img_shape=img_shape, + xml_path=xmlpath) \ No newline at end of file