YoloV11分割模型使用

官网

1. 官网文档：https://docs.ultralytics.com/zh
2. 代码仓库：https://github.com/ultralytics/ultralytics

使用分割模型预测

1. 安装依赖

 pip install ultralytics

2. 下载YOLO11x-seg模型：https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11x-seg.pt
3. 执行以下代码

import glob

import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.utils.plotting import Annotator, colors
import matplotlib

matplotlib.use('TkAgg')

import matplotlib.pyplot as plt  # noqa


class Runner:
    def __init__(self, model_path):
        self.model: YOLO = YOLO(model_path)
        # {idx: label}
        self.names: dict = self.model.model.names  # noqa

    def predict_and_show(self, img_or_path, conf=0.25):
        img = self.to_img(img_or_path)
        ori_img = img.copy()
        bg = np.zeros_like(img)
        results = self.model.predict(img, conf=conf)
        annotator = Annotator(img)

        for result in results:
            clss = result.boxes.cls.cpu().tolist()
            masks = result.masks.xy
            probs = result.masks.xyn

            for mask, prob, cls in zip(masks, probs, clss):
                color = colors(int(cls), True)
                txt_color = annotator.get_txt_color(color)
                label = self.names[int(cls)]
                annotator.seg_bbox(mask=mask, mask_color=color, label=label, txt_color=txt_color)
                self.fill_seg(bg, mask, mask_color=color, label=label, txt_color=txt_color)

        self.plt_imshow(ori_img, img, bg)

    @staticmethod
    def to_img(img_or_path):
        return cv2.imread(img_or_path) if isinstance(img_or_path, str) else img_or_path

    @staticmethod
    def plt_imshow(ori_img, label_img, mask_img):
        def _f(idx, img):
            plt.subplot(1, 3, idx)
            plt.imshow(img[:, :, [2, 1, 0]])  # cv2 bgr -> rgb
            plt.axis('off')

        plt.figure(figsize=(15, 6))
        _f(1, ori_img)
        _f(2, label_img)
        _f(3, mask_img)
        plt.show()

    @staticmethod
    def fill_seg(img, mask, mask_color, label=None, txt_color=(255, 255, 255)):
        if mask.size == 0:
            return

        lw = max(round(sum(img.shape) / 2 * 0.003), 2)
        tf = max(lw - 1, 1)  # font thickness
        sf = lw / 3  # font scale

        cv2.fillPoly(img, [np.int32(mask)], color=mask_color)
        text_size, _ = cv2.getTextSize(label, 0, sf, tf)

        if label:
            cv2.rectangle(
                img,
                (int(mask[0][0]) - text_size[0] // 2 - 10, int(mask[0][1]) - text_size[1] - 10),
                (int(mask[0][0]) + text_size[0] // 2 + 10, int(mask[0][1] + 10)),
                mask_color,
                -1,
            )
            cv2.putText(
                img, label, (int(mask[0][0]) - text_size[0] // 2, int(mask[0][1])), 0, sf, txt_color, tf
            )


def test():
    runner = Runner('yolo11x-seg.pt')  # 指定模型的路径

    for img_path in glob.glob(r'results\*.png'): # 指定图片的路径
        try:
            runner.predict_and_show(img_path)
        except AttributeError as e:
            print('==========>', e)


if __name__ == '__main__':
    test()

4. 结果显示
   

微调训练

标注图像

1. 安装labelme

pip install labelme

2. 启动labelme，命令行下执行：labelme
   
3. 然后选择打开目录，加载要标注的图像，再选择编辑 -> 创建多边形，就可以对图像进行标注了
   
4. 新版的labelme已经集成mask model功能，可根据鼠标的点击自动分割物体，我选的是EfficientSam(accuracy)，它会自动下载3个模型，默认放到当前用户的.cache/gdown下

efficient_sam_vits_encoder.onnx
https-COLON–SLASH–SLASH-github.com-SLASH-labelmeai-SLASH-efficient-sam-SLASH-releases-SLASH-download-SLASH-onnx-models-20231225-SLASH-efficient_sam_vits_decoder.onnx
https-COLON–SLASH–SLASH-github.com-SLASH-labelmeai-SLASH-efficient-sam-SLASH-releases-SLASH-download-SLASH-onnx-models-20231225-SLASH-efficient_sam_vits_encoder.onnx

使用时，从编辑的下拉列表中选择创建AI多边形。这时鼠标移动图像上，会变成一个十字形，在要标注的物体处点击，就会出现自动标注框，可以移动鼠标来区域选择，选好区域后，按Enter键，保存标注的类别名称

5. 可视化标注结果，在命令下执行

labelme_export_json.exe .\results\1-ori.json  # 指定标注生成的json文件路径

2025-01-23 16:58:21.247 | INFO | labelme.cli.export_json:main:66 - Saved to: .\results\1-ori

转成yolo要求的标注格式

import glob
import json
import os
from pathlib import Path
from typing import List, Optional

from tqdm import tqdm


def get_all_labels(json_dir_path: str) -> list:
    labels = list()
    for json_path in glob.glob(json_dir_path + '/*.json'):
        with open(json_path, 'r') as f:
            data = json.load(f)
            labels.extend([shape['label'] for shape in data['shapes']])
    return sorted(list(set(labels)))


def labelme_json_to_yolo_txt(json_dir_path, labels: Optional[List[str]] = None, save_dir_path: Optional[str] = None):
    """
    将labelme分割的json转成yolo训练时需要的数据格式
    """
    if labels is None:
        labels = get_all_labels(json_dir_path)
        
    print(f'labels: {labels}')

    for json_path in tqdm(glob.glob(json_dir_path + '/*.json')):
        with open(json_path, 'r') as load_f:
            json_dict = json.load(load_f)

        h, w = json_dict['imageHeight'], json_dict['imageWidth']

        txt_name = Path(json_path).name.replace('json', 'txt')
        if save_dir_path is None:
            txt_save_path = Path(json_dir_path).parent / txt_name
        else:
            os.makedirs(save_dir_path, exist_ok=True)
            txt_save_path = f'{save_dir_path}/{txt_name}'

        with open(txt_save_path, 'w') as f:
            for shape_dict in json_dict['shapes']:
                label = shape_dict['label']
                label_index = labels.index(label)
                points = shape_dict['points']

                points_nor_list = []

                for point in points:
                    points_nor_list.append(point[0] / w)
                    points_nor_list.append(point[1] / h)

                points_nor_list = list(map(lambda x: str(x), points_nor_list))
                points_nor_str = ' '.join(points_nor_list)

                label_str = str(label_index) + ' ' + points_nor_str + '\n'
                f.writelines(label_str)


if __name__ == "__main__":
    """
    使用labelme进行标注
    """
    labelme_json_to_yolo_txt(
        json_dir_path='data/images',  # 标注的图像和数据的目录
        save_dir_path='data/labels',  # 要保存的yolo数据格式目录
    )

分割数据集

将图像和标注数据分割成训练集和验证集

import glob
import os
import random
import shutil
from pathlib import Path


def copy(source_img_dir_path, txt_path_ls, start_idx, end_idx, target_label_dir_path, target_img_dir_path):
    for txt_path in txt_path_ls[start_idx: end_idx]:
        img_path = None
        txt_name = Path(txt_path).name

        for suffix in ['.jpg', '.png', '.bmp', '.jpeg']:
            img_name = txt_name.split('.')[0] + suffix
            if Path(f'{source_img_dir_path}/{img_name}').exists():
                img_path = f'{source_img_dir_path}/{img_name}'
                break
        if img_path is None:
            print(f'xxx> {txt_name}的图像不存在')
            continue

        shutil.copyfile(txt_path, target_label_dir_path + '/' + txt_name)
        shutil.copyfile(img_path, target_img_dir_path + '/' + img_name)


def main(img_dir_path, txt_dir_path, save_dir_path, train_percent = 0.85):
    os.makedirs(save_dir_path, exist_ok=True)

    images_dir_path = os.path.join(save_dir_path, 'images')
    labels_dir_path = os.path.join(save_dir_path, 'labels')

    img_train_path = os.path.join(images_dir_path, 'train')
    img_val_path = os.path.join(images_dir_path, 'val')

    label_train_path = os.path.join(labels_dir_path, 'train')
    label_val_path = os.path.join(labels_dir_path, 'val')

    [
        os.makedirs(path, exist_ok=True)
        for path in [
            images_dir_path, labels_dir_path,
            img_train_path, label_train_path,
            img_val_path, label_val_path
        ]
    ]

    txt_path_ls = glob.glob(txt_dir_path + '/*.txt')
    random.shuffle(txt_path_ls)
    train_num = int(len(txt_path_ls) * train_percent)
    print(f"训练集数目：{train_num}, 验证集数目：{len(txt_path_ls) - train_num}")

    copy(img_dir_path, txt_path_ls, 0, train_num, label_train_path, img_train_path)
    copy(img_dir_path, txt_path_ls, train_num, len(txt_path_ls), label_val_path, img_val_path)


if __name__ == '__main__':
    main(
        'data/images',  # 图像的目录
        'data/labels',  # 转成yolo的标注数据目录
        'data/dataset', # 保存的目录
        train_percent=0.95
    )

启动训练

1. 将yolo11x-seg.pt放在当前目录下
2. 下载yolo11n.pt模型，放在当前目录下。也可不提前下载，代码会自动下载。
3. 在当前目录下，创建训练配置文件：seg-train.yaml

train: data/dataset/images/train
val: data/dataset/images/val
nc: 10
names: ['bear', 'book', 'cell-phone', 'chess', 'cola', 'cup', 'keyboard', 'mouse', 'plant', 'tv']

4. 在当前目录下，创建python文件，并执行

from ultralytics import YOLO


if __name__ == '__main__':
    model = YOLO('yolo11x-seg.pt') 
    model.train(data='seg-train.yaml', epochs=200)

训练完成后：