YOLOV5_7.0 数据集自动标注

YOLOV5_7.0 数据集自动标注，只适用于7.0版本

1. 创建文件夹

在桌面创建文件夹VOCData

VOCData内创建两个文件夹images和Annotations分步别用来存放照片和自动标注生成的xml标签文件
在detect.py同级目录下新建Auto_Lab.py文件

Auto_Lab.py文件核心代码：

import sys

from utils.dataloaders import *
from utils.general import check_img_size, non_max_suppression, scale_boxes
from utils.torch_utils import select_device #, time_synchronized
import argparse
import os
import shutil
from models.experimental import attempt_load
import cv2
import torch
import torch.backends.cudnn as cudnn
import numpy as np

from xml.etree import ElementTree as ET
import warnings

warnings.filterwarnings('ignore')


# 定义一个创建一级分支object的函数
def create_object(root, xi, yi, xa, ya, obj_name):  # 参数依次，树根，xmin，ymin，xmax，ymax
    # 创建一级分支object
    _object = ET.SubElement(root, 'object')
    # 创建二级分支
    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'
    # 创建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


# 创建xml文件的函数
def create_tree(sources, image_name, h, w):
    imgdir = sources.split('/')[-1]
    # 创建树根annotation
    annotation = ET.Element('annotation')
    # 创建一级分支folder
    folder = ET.SubElement(annotation, 'folder')
    # 添加folder标签内容
    folder.text = (imgdir)

    # 创建一级分支filename
    filename = ET.SubElement(annotation, 'filename')
    filename.text = image_name

    # 创建一级分支path
    path = ET.SubElement(annotation, 'path')

    path.text = '{}/{}'.format(sources, image_name)  # 用于返回当前工作目录

    # 创建一级分支source
    source = ET.SubElement(annotation, 'source')
    # 创建source下的二级分支database
    database = ET.SubElement(source, 'database')
    database.text = 'Unknown'

    # 创建一级分支size
    size = ET.SubElement(annotation, 'size')
    # 创建size下的二级分支图像的宽、高及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'

    # 创建一级分支segmented
    segmented = ET.SubElement(annotation, 'segmented')
    segmented.text = '0'
    return annotation


def detect(opt, model, img, img0):
    device = select_device(opt.device)
    half = device.type != 'cpu'  # half precision only supported on CUDA

    # img = img_transpose(img0, imgsz, 32)
    img = torch.from_numpy(img).to(device)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    if img.ndimension() == 3:
        img = img.unsqueeze(0)

    pred = model(img, augment=opt.augment)[0]

    # Apply NMS
    pred = non_max_suppression(
        pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
    # Process detections
    for i, det in enumerate(pred):  # detections per image
        if det is not None and len(det):
            # Rescale boxes from img_size to im0 size
            det[:, :4] = scale_boxes(img.shape[2:], det[:, :4], img0.shape).round()
    return det


def main(opt):
    source, weights, imgsz, = opt.source, opt.weights, opt.img_size

    # Initialize
    device = select_device(opt.device)

    half = device.type != 'cpu'  # half precision only supported on CUDA

    # Load model
    model = attempt_load(weights, device=None, inplace=True, fuse=True)  # load FP32 model
    stride = int(model.stride.max())  # model stride
    imgsz = check_img_size(imgsz, s=stride)  # check img_size
    if half:
        model.half()  # to FP16

    img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img
    # run once
    _ = model(img.half() if half else img) if device.type != 'cpu' else None
    # Get names and colors
    names = model.module.names if hasattr(model, 'module') else model.names

    dataset = LoadImages(source, img_size=imgsz, stride=stride)
    print(dataset)

    # images_list = os.listdir(source)
    # images_style = ['.jpg', '.png', '.bmp']
    # images_list = [x for x in images_list if x[-4:] in images_style]
    # print(images_list)

    for path, img, im0s, cap, vid_cap in dataset:
        image_name = os.path.split(path)[-1]
        # print('path:', path)

        # 检测饲料袋
        boxes = detect(opt, model, img, im0s)
        # print(len(boxes))
        (h, w) = im0s.shape[:2]
        annotation = create_tree(source, image_name, h, w)
        # print(annotation)

        for box in boxes:
            if float(box[4]) > opt.conf_thres:
                x1, y1, x2, y2, label_id = int(box[0]), int(box[1]), int(box[2]), int(box[3]), int(box[5])
                label = names[int(label_id)]
                # print(x1, y1, x2, y2, label)
                create_object(annotation, x1, y1, x2, y2, label)

        tree = ET.ElementTree(annotation)
        annotation_path_root = source.replace(source.split('/')[-1], 'annotations')
        tree.write('{}/{}.xml'.format(annotation_path_root, image_name[:-4]))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--weights', type=str,
                        default=r'E:/Python_Learn_Source/YOLO/yolov5_7.0_Attention_Multiple/runs/train/yolov5s/weights/best.pt', help='model.pt path')
    # file/folder, 0 for webcam
    parser.add_argument('--source', type=str,
                        default=r'C:/Users/17616/Desktop/VOCData/images', help='source')
    parser.add_argument('--output', type=str, default='inference/output',
                        help='output folder')  # output folder
    parser.add_argument('--img-size', type=int, default=640,
                        help='inference size (pixels)')
    parser.add_argument('--conf-thres', type=float,
                        default=0.25, help='object confidence threshold')
    parser.add_argument('--iou-thres', type=float,
                        default=0.05, help='IOU threshold for NMS')
    parser.add_argument('--device', default='cpu',
                        help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--classes', nargs='+', type=int,
                        help='filter by class')
    parser.add_argument('--augment', action='store_true',
                        help='augmented inference')
    parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
    args = parser.parse_args()
    args.img_size = check_img_size(args.img_size)
    print(args)
    with torch.no_grad():
        main(args)

需要修改下图所示的两个路径：

修改完运行Auto_Lab.py文件

2.xml转txt、数据集划分

参考yolov5 7.0 + 划分/制作自己数据