使用YOLOv8实现YOLO格式标签的自动标注

1 准备

（1）已经训练好的yolov8权重文件（'yolov8/runs/detect/train/weights/best.pt'）

（2）需要自动标注的图片文件（'/image/*.jpg'）

2 开始标注

import torch
from ultralytics import YOLO


class Yolov8Detect():
    def __init__(self, weights):
        cuda = True if torch.cuda.is_available() else False
        self.device = torch.device('cuda:0' if cuda else 'cpu')
        self.detect_model = YOLO(weights)
        self.detect_model.to(self.device)


    def inferences(self, inputs):
        results = self.detect_model(inputs)
        for result in results:
            label_text = []
            boxes = result.boxes
            for box in boxes:
                cat_num = int(box.cls.cpu())
                label_text.append([cat_num, box.xywhn.cpu().numpy()])
            save_path = inputs.replace('jpg', 'txt')
            txt_construct(save_path, label_text=label_text)


def txt_construct(save_path, label_text):
    with open(save_path, 'w') as file:
        file.truncate()
    for label in label_text:
        with open(save_path, 'a') as txt_file:
            label_ = label[0]
            size = label[1][0].tolist()
            size_string = ' '.join(map(str, size))
            result = f'{label_} {size_string}'
            print('result', result)
            txt_file.write(str(result))
            txt_file.write('\n')



if __name__ == '__main__':

    #model_path替换成自己的model路径
    model_path = '/best.pt'
    model = Yolov8Detect(model_path)
    import glob

    #image_path替换成需要自动标注的文件路径
    image_path = glob.glob('/*.jpg')
    for img_path in image_path[:]:
        model.inferences(img_path)

  生成的.txt文件会保存在图片所在文件夹中。