Ultralytics 团队付出了巨大的努力，使创建自定义 YOLO 模型变得非常容易。但是，处理大型数据集仍然很痛苦。训练 yolo 分割模型需要数据集具有其特定格式，这可能与你从大型数据集中获得的格式不完全相同。如果你想使用巨大的 OpenImagesV7 作为图像和标签的来源，情况就是如此。

在本教程中，我们将介绍如何从 OpenImagesV7 获取数据（图像和分割掩码）；如何将其转换为 YOLO 格式（这是本教程中最复杂的部分）；以及如何使用我们的数据集训练 yolov8-seg 模型的预览。

NSDT工具推荐： Three.js AI纹理开发包 - YOLO合成数据生成器 - GLTF/GLB在线编辑 - 3D模型格式在线转换 - 可编程3D场景编辑器 - REVIT导出3D模型插件 - 3D模型语义搜索引擎 - AI模型在线查看 - Three.js虚拟轴心开发包 - 3D模型在线减面 - STL模型在线切割

1、环境

要清楚：本教程需要 Python 3（在 3.10 下测试）。作为基础图像，我使用 AWS Sagemaker conda_pytorch_p310 ，它包含 PyTorch 和许多常用工具，如 Numpy 和 OpenCV。不过，我们需要一些额外的包来覆盖它：

sudo yum install -y openssl-devel openssl11-libs libcurl
pip install --upgrade pip setuptools wheel
pip install fiftyone
pip install fiftyone-db-rhel7 --force-reinstall
pip install shapely polars
pip install ultralytics

2、数据集

我们将使用 Google OpenImages Dataset v7 来训练我们的模型。这个数据集非常庞大，有数百万张图像，旨在完成一系列计算机视觉任务，例如对象检测、分类和实例分割。因此，每张图像都与这些任务的标签配对，涵盖了最常见的多种对象：人、脸、狗、猫、汽车、树木等。请注意，并非所有标签都适用于所有任务，因此在数据集网站上进行一些探索以了解它涵盖了您正在尝试构建的模型的扩展范围是值得的。对于大多数计算机视觉问题来说，这可能已经足够了。

因此，要实际获取数据，有一个很棒的工具叫做 FiftyOne，它能够只下载你需要的数据。这避免了下载和处理整个数据集的负担。我是否已经提到它非常庞大？如果没有 FiftyOne，处理起来真的很难。

第一步是选择一个好的位置来下载数据集。我假设 Sagemaker 环境采用默认设置，但您可以选择更适合您的设置。请记住，在此路径中将自动创建文件夹 open-images-v7，以保存数据。

# choose your prefered path to download the dataset
# a folder named open-images-v7 will be created automatically inside of it
dataset_path = '/home/ec2-user/SageMaker/dataset'

import os
import torch
import torchvision

import fiftyone as fo
fo.config.default_ml_backend = "torch"
fo.config.dataset_zoo_dir = dataset_path

下一个函数只是为了让我们的脚本更简洁，也更灵活。我们可以简单地将所有想要的标签放入 fiftyone 下载器函数中，让它决定要获取每个标签的样本数量。但是，在这里我需要更好地控制要下载每个类别的样本数量。当然，如果确实有可用的样本， fiftyone 只会下载我们要求的数字，否则它会尽可能多地获取样本。

def download_dataset(split, classes, max_samples=None):
    print(f'>> Split: {split}, classes: {classes}, max_samples: {max_samples}')
    return fo.zoo.load_zoo_dataset(
            "open-images-v7",
            label_types=["segmentations"],
            drop_existing_dataset=False,
            split=split,
            classes=classes,
            max_samples=max_samples,
    )

在调用下载器函数之前，我们先描述一下我们需要的标签以及每个标签的数量。请注意，如果我们输入 None，我们就可以实现“尽可能多”的行为，而无需猜测一个很大的随机数。为了简单起见，我们将获取包含人或汽车的样本。最多一千个人物样本和尽可能多的汽车样本。我们希望 70% 用于训练，20% 用于验证，10% 用于测试；这些百分比不能保证，因为每个子集可能都没有足够的样本。

target_split = {'train': 0.7, 'validation': 0.2, 'test': 0.1}
target_classes = {
    "Person": 1_000,
    "Car": None
}

现在我们可以运行下载器函数了。为此，我们将遍历上面定义的目标类。

for cls_name, total in target_classes.items():
    for split_name, split_pct in target_split.items():
        max_samples = int(total * split_pct) if total is not None else None
        download_dataset(split=split_name, classes=[cls_name], max_samples=max_samples)

即使在强大的 AWS Sagemaker 环境中，脚本也可能需要一段时间才能完成。当然，时间与请求的样本数量成正比。下载数据集后，我们仍未准备好将其与 yolo 一起使用。数据以彩色图像的形式出现，并配有一组分割蒙版（灰度图像），图像中的每个对象都有一个蒙版。但是，yolo 需要将彩色图像与文本文件配对，文本文件的每一行都描述对象的类别和描述其多边形的坐标。请查看 yolo 文档。这将引出我们的下一章，介绍如何将数据集转换为适合 YoloV8 的格式。

3、将 OpenImagesV7 转换为 Yolo 分割

原始分割标签以灰度图像的形式提供。如前所述，Yolo 要求分割标签位于文本文件中，其中包含图像中每个对象的一行，遵循以下模式：对象类 ID，然后是描述对象多边形的 XY 列表。

这种转换需要一些工具来加载图像、将蒙版转换为多边形、降低多边形复杂性并将多边形写在文本文件中。所以让我们在一个新的干净的笔记本或 Python 脚本中一步一步地完成这项工作。首先是基本配置：

# base path must be the same as our previous dataset_path
base_path = '/home/ec2-user/SageMaker/dataset'

# destination of the converted dataset
target_path = '/home/ec2-user/SageMaker/dataset-yolo'

# a list with same keys as on fetch.py
target_classes = [
  "Person",
  "Car"
]

import os
import cv2
import yaml
import shutil
import pandas as pd
import polars as pl
import multiprocessing
import numpy as np
from tqdm import tqdm
from joblib import Parallel, delayed
from shapely.geometry import Polygon
from matplotlib import pyplot as plt

现在我们定义一些函数来执行从灰度蒙版图像到简化多边形，然后再到 XY 列表的实际转换：

### Mask to Poly ###

def mask_to_polygon(mask_path):
    mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
    contours, _ = cv2.findContours(
        mask,
        cv2.RETR_EXTERNAL,
        cv2.CHAIN_APPROX_SIMPLE
    )

    polygons = []
    for contour in contours:
        polygon = contour.reshape(-1, 2)
        polygon_norm = polygon.astype(float)
        polygon_norm[:, 0] /= mask.shape[1]  # X
        polygon_norm[:, 1] /= mask.shape[0]  # Y
        polygon_norm = np.round(polygon_norm, 4)

        polygon_shapely = Polygon(polygon_norm)
        polygon_simplified = polygon_shapely.simplify(0.002, preserve_topology=True)
        polygons.append(polygon_simplified)

    return polygons

def polygon_to_yolo(polygon):
    x, y = polygon.exterior.coords.xy
    xy = []
    for xx, yy in zip(x, y):
        xy.append(xx)
        xy.append(yy)
    return xy

def polygon_to_mask(polygon, shape):
    mk = np.zeros(shape, dtype=np.uint8)
    x, y = polygon.exterior.coords.xy
    xy = [
        [int(xx * shape[1]), int(yy * shape[0])]
        for xx, yy in zip(x, y)
    ]
    cv2.fillConvexPoly(mk, np.array(xy, dtype='int32'), color=255)
    return mk

准备好图像处理工具后，下一步就是加载标签位置。这将遍历数据库文件系统查看图像文件，对于每幅图像，它将查看分割索引并获取相应的蒙版图像位置。到目前为止，图像仍未打开。有趣的事实：这里使用的是 polars 而不是 pandas，因为这个简单任务的处理速度差异很大。

### loading openimagesv7 labels ###

class_list_filepath = os.path.join(base_path, 'train/metadata/classes.csv')
class_df = pd.read_csv(class_list_filepath, header=None, names=['URI', 'ClassName'])
class_map_r = dict(zip(class_df.URI, class_df.ClassName))
class_map_r = {k: v for k, v in class_map_r.items() if v in target_classes}

# convert from openimagev7 label hash to an integer
class_map = { k: i for i, k in enumerate(list(class_map_r.keys()))}
# class_map = {
# '/m/01g317':  0, # 'Person'
# '/m/0k4j':    1, # 'Car'
# }
print('class_map:')
print(class_map)

def get_image_file_names(directory):
    image_extensions = ['.jpg', '.jpeg', '.png', '.gif', '.bmp']  # Add more extensions if needed
    image_file_names = set()

    for filename in os.listdir(directory):
        nm, ext = os.path.splitext(filename)
        if ext in image_extensions:
            image_file_names.add(nm)

    return image_file_names
  
def load_labels(split_name):
    df = pl.read_csv(os.path.join(base_path, split_name, 'labels/segmentations.csv'))
    df = df[['MaskPath', 'ImageID', 'LabelName']]

    image_ids = get_image_file_names(os.path.join(base_path, split_name, 'data'))
    df = df.filter(pl.col('ImageID').is_in(image_ids))
    
    target_ids = set(class_map.keys())
    df = df.filter(pl.col('LabelName').is_in(target_ids))

    df = df.with_columns(pl.col('MaskPath').map_elements(lambda x: x[0].upper()).alias('Subdir'))
    df = df.with_columns((base_path + f'{split_name}/labels/masks/' + pl.col('Subdir') + '/' + pl.col('MaskPath')).alias('MaskFullPath'))
    
    df = df.with_columns(pl.col(['LabelName']).map_dict(class_map).alias('LabelID'))

    return df


train_df = load_labels('train')    
valid_df = load_labels('validation')
test_df = load_labels('test')

最后，是时候遍历所有掩模图像，将其转换为 XY 多边形列表并按照 yolo 标准将结果写入文本文件中。

def macro_mask2yolopoly(p):
    try:
        poly = mask_to_polygon(p)
        xy = polygon_to_yolo(poly[0])
        return xy
    except:
        return []
    return []


def conv_mask_xy(df):
    return df.with_columns(
        pl.col('MaskFullPath').map_elements(
            lambda p: macro_mask2yolopoly(p)
        ).alias('XY')
    )

train_df = conv_mask_xy(train_df)
valid_df = conv_mask_xy(valid_df)
test_df = conv_mask_xy(test_df)

def write_yolo_labels(df, subset, persistence=True):
    df = df.filter(pl.col('XY').map_elements(len) > 0)
    
    df = df.with_columns(
        pl.col('XY').map_elements(lambda xy: xy.map_elements(lambda e: str(e))).list.join(' ').alias('TXY'))
    df = df.with_columns(
        (pl.col('LabelID').cast(pl.Utf8) + ' ' + pl.col('TXY')).alias('Sample'))
    
    g = df.group_by('ImageID').agg(['Sample'])
    g = g.with_columns(pl.col('Sample').list.join('\n').alias('StrSamples'))
    g = g.with_columns((target_path + subset + '/' + pl.col('ImageID') + '.txt').alias('Path'))

    if persistence:
        os.makedirs(os.path.join(target_path, subset), exist_ok=True)
        for row in g.iter_rows(named=True):
            with open(row['Path'], 'w') as f:
                f.write(row['StrSamples'])

    return g
            
train_df = write_yolo_labels(train_df, 'train')
valid_df = write_yolo_labels(valid_df, 'validation')
test_df = write_yolo_labels(test_df, 'test')

好吧，我们的新数据集最终被拆分成不同位置的图像和 yolo 标签。因此，为了统一，我们需要将下载的图像连同我们的标签一起复制到目标位置。这里有几种可能性，比如在 openimagesv7 文件夹中创建标签、移动图像等。我更喜欢复制并保留原始文件，以防脚本失败，稍后我可以简单地删除整个 openimages 缓存以释放一些空间。

def copy_data(df, subset):
    for iid in df.select(pl.col('ImageID')).get_columns()[0].to_list():
        try:
            fnm = f"{iid}.jpg"
            src = os.path.join(base_path, subset, "data", fnm)
            dst = os.path.join(target_path, subset)
            # print(f'{src} -> {dst}')
            shutil.copy2(src, dst)
        except:
            continue

copy_data(valid_df, 'validation')
copy_data(test_df, 'test')
copy_data(train_df, 'train')

Yolo 数据集仅包含描述它的 YAML 文件才完整。创建一个很容易，就像这样：

from pathlib import Path

yaml_content = f'''
path: /home/ec2-user/SageMaker/dataset-yolo
train: train
val: validation
test: test

# Classes - use the class_map as guide
names:
  0: person
  1: car
'''

with Path(os.path.join(target_path, 'seg_dataset.yaml')).open('w') as f:
    f.write(yaml_content)

4、训练 YoloV8-Seg 模型

本教程最难的部分是让数据集适合 Yolo。鉴于有更好的关于如何训练 Yolo 模型的教程，我将简短地讲解本课程。还请查看 Ultralytics 文档。但为了让你不要抱怨这是一个半生不熟的教程，我们开始吧：

# on the command line (CLI)
yolo segment train data=/home/ec2-user/SageMaker/dataset-yolo/seg_dataset.yaml model=yolov8n-seg.pt epochs=100 imgsz=640

这将使用预先训练好的 yolov8n-seg.pt 作为我们自定义分割模型的基础，并将更新超过 100 个 epoch。模型的输出可能位于 ./runs/exp/ 上。如果不是这种情况，请查看文档。

5、参考文档

• vchaparro 的代码：将掩码转换为多边形
• Alon Lekhtman 关于训练 YoloV8-Seg 的 Medium 帖子
• Google OpenImagesV7 数据集
• Ultralytics

原文链接：分割掩码转YOLO格式 - BimAnt