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

F.M.H

已于 2025-05-22 13:02:41 修改

阅读量1.5k

点赞数 11

CC 4.0 BY-SA版权

文章标签： YOLO 人工智能深度学习

于 2024-01-09 10:49:18 首次发布

本文链接：https://blog.csdn.net/weixin_43491146/article/details/135474446

本文介绍了如何使用预训练的YOLOv8模型对图片进行自动标注，包括加载模型权重、进行推理并保存标注结果到.txt文件的过程。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

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文件会保存在图片所在文件夹中。