Automatic annotation of target detection based on lffd model to generate XML file

import os
from os import getcwd
from xml.etree import ElementTree as ET
# from lxml import etree as ET
import argparse
import sys
import cv2
from vision.ssd.config.fd_config import define_img_size
from vision.ssd.mb_tiny_fd import create_mb_tiny_fd, create_mb_tiny_fd_predictor
from vision.ssd.mb_tiny_RFB_fd import create_Mb_Tiny_RFB_fd, create_Mb_Tiny_RFB_fd_predictor


#  Define a to create a first level branch object Function of 
def create_object(root, xi, yi, xa, ya, obj_name):  #  The parameters are in turn , The root ,xmin,ymin,xmax,ymax
    #  Create a first level branch object
    _object = ET.SubElement(root, 'object')
    #  Create a secondary branch 
    name = ET.SubElement(_object, 'name')
    # print(obj_name)
    name.text = str(obj_name)
    pose = ET.SubElement(_object, 'pose')
    pose.text = 'Unspecified'
    truncated = ET.SubElement(_object, 'truncated')
    truncated.text = '0'
    difficult = ET.SubElement(_object, 'difficult')
    difficult.text = '0'
    #  establish bndbox
    bndbox = ET.SubElement(_object, 'bndbox')
    xmin = ET.SubElement(bndbox, 'xmin')
    xmin.text = '%s' % xi
    ymin = ET.SubElement(bndbox, 'ymin')
    ymin.text = '%s' % yi
    xmax = ET.SubElement(bndbox, 'xmax')
    xmax.text = '%s' % xa
    ymax = ET.SubElement(bndbox, 'ymax')
    ymax.text = '%s' % ya


#  establish xml Functions of files 
def create_tree(image_name, h, w, imgdir):
    global annotation
    #  Create tree roots annotation
    annotation = ET.Element('annotation')
    #  Create a first level branch folder
    folder = ET.SubElement(annotation, 'folder')
    #  add to folder Label content 
    folder.text = (imgdir)

    #  Create a first level branch filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = image_name

    #  Create a first level branch path
    path = ET.SubElement(annotation, 'path')

    path.text = getcwd() + '/{}/{}'.format(imgdir, image_name)  #  Used to return to the current working directory 

    #  Create a first level branch source
    source = ET.SubElement(annotation, 'source')
    #  establish source The second branch under database
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'

    #  Create a first level branch size
    size = ET.SubElement(annotation, 'size')
    #  establish size Width of secondary branch image under 、 Gao Ji depth
    width = ET.SubElement(size, 'width')
    width.text = str(w)
    height = ET.SubElement(size, 'height')
    height.text = str(h)
    depth = ET.SubElement(size, 'depth')
    depth.text = '3'

    #  Create a first level branch segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'

#  be used for xml Format indent 
def pretty_xml(element, indent, newline, level=0):  # elemnt For what came in Elment class , Parameters indent Used to indent ,newline For line breaks 
    if element:  #  Judge element Whether there are child elements 
        if (element.text is None) or element.text.isspace():  #  If element Of text No content 
            element.text = newline + indent * (level + 1)
        else:
            element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * (level + 1)
            # else: #  If you remove the comments in these two lines ,Element Of text There will be another line 
            # element.text = newline + indent * (level + 1) + element.text.strip() + newline + indent * level
    temp = list(element)  #  take element Turn into list
    for subelement in temp:
        if temp.index(subelement) < (len(temp) - 1):  #  If not list Last element of , Indicates that the next line is the beginning of the same level element , Indent should be consistent 
            subelement.tail = newline + indent * (level + 1)
        else:  #  If it is list Last element of ,  Indicates that the next line is the end of the parent element , The indent should be one less 
            subelement.tail = newline + indent * level
        pretty_xml(subelement, indent, newline, level=level + 1)  #  Recursive operation on child elements 


if __name__ == '__main__':
    imgdir =  "ext_imgs"
    outdir = "./annotation_generator"
    if not os.path.exists(outdir):
        os.makedirs(outdir)
    label_path = "./models/voc-model-labels.txt"

    model_path = "models/train-version-RFB/RFB-Epoch-199-Loss-1.9399930238723755.pth"

    parser = argparse.ArgumentParser(
        description='ext_imgs')

    parser.add_argument('--net_type', default="RFB", type=str,
                        help='The network architecture ,optional: RFB (higher precision) or slim (faster)')
    parser.add_argument('--input_size', default=640, type=int,
                        help='define network input size,default optional value 128/160/320/480/640/1280')
    parser.add_argument('--threshold', default=0.6, type=float,
                        help='score threshold')
    parser.add_argument('--candidate_size', default=1500, type=int,
                        help='nms candidate size')
    parser.add_argument('--path', default="ext_imgs", type=str,
                        help='imgs dir')
    parser.add_argument('--test_device', default="cuda:0", type=str,
                        help='cuda:0 or cpu')
    args = parser.parse_args()
    define_img_size(args.input_size)
    test_device = args.test_device
    class_names = [name.strip() for name in open(label_path).readlines()]
    net = create_Mb_Tiny_RFB_fd(len(class_names), is_test=True, device=test_device)
    predictor = create_Mb_Tiny_RFB_fd_predictor(net, candidate_size=args.candidate_size, device=test_device)

    net.load(model_path)

    IMAGES_LIST = os.listdir(imgdir)

    for image_name in IMAGES_LIST:
        #print(image_name)
        if image_name.endswith(".jpg"):
            image = cv2.imread(os.path.join(imgdir, image_name))
            images = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            (h, w) = image.shape[:2]
            create_tree(image_name, h, w, imgdir)

            boxes, labels, probs = predictor.predict(images, args.candidate_size / 2, args.threshold)
            # print(boxes)
        for i in range(boxes.size(0)):
            box = boxes[i, :]
            # create_object(annotation, int(box[0]), int(box[1]), int(box[2]), int(box[3]), image)
            create_object(annotation, int(box[0]), int(box[1]), int(box[2]), int(box[3]), "ext")

        tree = ET.ElementTree(annotation)
        root = tree.getroot()
        pretty_xml(root, '\t', '\n')
        tree.write('./{}/{}.xml'.format(outdir, image_name.strip('.jpg')), encoding='utf-8')
        print(image_name)