新添加自动标注部分

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  # 二级目录的值 #结果展示：<id>1</id>
+            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)
+
+

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()
+ 
+
+
+
+

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)