基于YOLOv8的目标检测半自动标注脚本

YOLOv8安装、训练见这篇文章【YOLOv8教程】（手把手详细版，支持训练自定义数据集），包含环境配置、数据配置、训练、预测、验证、导出

1.本项目文件结构如下：

├── README.md
├── README.zh-CN.md
├── class_index.py
├── config.py
├── create_xml.py
├── docker
├── docs
├── examples
├── mkdocs.yml
├── requirements.txt
├── setup.cfg
├── setup.py
├── tests
├── ultralytics

2.目标检测自动化标注脚本

只需要一个脚本文件即可完成费时费力的标注工作
创建creat_xml.py，创建位置见上述目录，其中内容如下：

from lxml.etree import Element, tostring, parse
from lxml.etree import SubElement as subElement
import torch


from ultralytics import YOLO


class Yolov8Detect():
    def __init__(self, weights, img_size=960, conf_thresh=0.2, iou_thresh=0.5, batchsize=16):
        cuda = True if torch.cuda.is_available() else False
        self.device = torch.device('cuda:0' if cuda else 'cpu')
        self.detect_model = YOLO(weights)
        # print(weights)
        self.detect_model.to(self.device)

        #####
        # 0.x
        # self.img_size = img_size
        # self.batch_size = batchsize
        # self.conf_thresh = conf_thresh
        # self.iou_thresh = iou_thresh
        # self.names =['ganjin','guahua','liewen','heidian','liangshao','jingshangfen','gouya','juzui','gehen','lvgai','yabutie']
        self.names = ['ganjing', 'guahua', 'liewen_shen', 'heidian', 'yaofenxian', 'jingshangfen', 'gouya', 'juzui',
                      'gehen', 'lvgai', 'yabutie', 'gaisong', 'quegai', 'pingzi', 'liewen_di','bianxing', 'liang','aoxian','xianglianghuan','ganjin_fx', 'boli']
        
    def inferences(self, inputs):
        # result1 = self.detect_model.predict()
        results = self.detect_model(inputs)
        for result in results:
            label_text = []
            boxes = result.boxes
            for box in boxes:
                cat_num = int(box.cls.cpu())
                cate=self.names[cat_num]
                label_text.append([cate, box.xyxy])
        image_name = inputs.split('/')[-1]
        save_path = inputs.replace('jpg', 'xml')
        xml_construct(save_path, inputs.split('/')[-2], image_name, inputs, width=1920, height=1200, label_text=label_text)
        print('save_path:', save_path)
                
      

def xml_construct(save_path,folder,filename,path,label_text,width=800,height=600,depth = 3,segmented=0):
    default_text = 'default'
    node_root = Element('annotation')  # 根节点
 
    node_folder = subElement(node_root, 'folder')  # 在节点下添加名为'folder'的子节点
    node_folder.text = folder  # 设定节点的文字
 
    node_filename = subElement(node_root, 'filename')
    node_filename.text = filename
 
    node_path = subElement(node_root, 'path')
    node_path.text = path
 
    node_size = subElement(node_root, 'size')
    node_size_width = subElement(node_size, 'width')
    node_size_width.text = '%s' % int(width)
    node_size_height = subElement(node_size, 'height')
    node_size_height.text = '%s' % int(height)
    node_size_depth = subElement(node_size, 'depth')
    node_size_depth.text = '%s' % int(depth)
 
    node_segmented = subElement(node_root, 'segmented')
    node_segmented.text = '%s' % int(segmented)
    
    for label in label_text:
        
        node_size = subElement(node_root, 'object')
        node_size_width = subElement(node_size, 'name')
        node_size_width.text = '%s' % (label[0])
        node_size_height = subElement(node_size, 'pose')
        node_size_height.text = 'Unspecified' 
        node_size_depth = subElement(node_size, 'truncated')
        node_size_depth.text = '0' 
        node_size_depth = subElement(node_size, 'difficult')
        node_size_depth.text = '0' 
        
        node_bndbox = subElement(node_size, 'bndbox')
        node_bndbox_xmin = subElement(node_bndbox, 'xmin')
        node_bndbox_xmin.text = '%s' % int(label[1][0][0])
        node_bndbox_ymin = subElement(node_bndbox, 'ymin')
        node_bndbox_ymin.text = '%s' % int(label[1][0][1])
        node_bndbox_xmax = subElement(node_bndbox, 'xmax')
        node_bndbox_xmax.text = '%s' % int(label[1][0][2])
        node_bndbox_ymax = subElement(node_bndbox, 'ymax')
        node_bndbox_ymax.text = '%s' % int(label[1][0][3])    
    
 
    xml = tostring(node_root, pretty_print=True) #将上面设定的一串节点信息导出
    with open(save_path,'wb') as f: #将节点信息写入到文件路径save_path中
        f.write(xml)

    return
 

if __name__ == '__main__':
    model_path = '/home/wanglu/runs/detect/train11_gray/weights/best.pt'
    model = Yolov8Detect(model_path)
    import glob 
    image_path = glob.glob('/home/wanglu/detection/0524/liewen_gray/*.jpg')
    for img_path in image_path[:]:
        model.inferences(img_path)

需要更改的地方：
model_path 和image_path更改为自己的模型和图片路径即可

脚本运行后，xml文件会生成在图片路径下。

3.如果想将xml和图像文件分离，可使用如下脚本:

import os
import shutil

path_xml = "/home/wanglu/detection/guahua_0607v2_gray/"
filelist = os.listdir(path_xml)
# path1 = "/home/wanglu/detection/guahua_0607v2/"
# path2 = "/home/wanglu/detection/one/1jiweilabels/"
path3 = "/home/wanglu/detection/guahua_0607v2_gray_xml/"


for files in filelist:
    filename1 = os.path.splitext(files)[1]  # 读取文件后缀名
    filename0 = os.path.splitext(files)[0]  #读取文件名
    # print(filename1)
    # m = filename1 == '.txt'
    # print(m)
    # if filename1 == '.txt' :
    #     full_path = os.path.join(path_xml, files)
    #     despath = path2 + filename0 +'.txt' #.txt为你的文件类型，即后缀名，读者自行修改
    #     shutil.move(full_path, despath)

    # if filename1 == '.jpg':
    #     full_path = os.path.join(path_xml, files)
    #     despath = path1 + filename0 + '.jpg'  #.jpg为你的文件类型，即后缀名，读者自行修改
    #     shutil.move(full_path, despath)

    if filename1 == '.xml':
        full_path = os.path.join(path_xml, files)
        despath = path3 + filename0 + '.xml'  # .xml为你的文件类型，即后缀名，读者自行修改
        shutil.move(full_path, despath)