步骤一：构建YOLO数据集#

由于比赛原始数据集较大，我们采样部分数据构建训练集和验证集：

1
if os.path.exists('yolo_seg_dataset'):
2
    shutil.rmtree('yolo_seg_dataset')
3

4
os.makedirs('yolo_seg_dataset/train')
5
os.makedirs('yolo_seg_dataset/valid')
6

7
def normalize_polygon(polygon, img_width, img_height):
8
    return [(x / img_width, y / img_height) for x, y in polygon]
9

10
# 采样训练集
11
for row in training_anno.iloc[:10000].iterrows():
12
    shutil.copy(row[1].Path, 'yolo_seg_dataset/train')
13

14
    img = cv2.imread(row[1].Path)
15
    img_height, img_width = img.shape[:2]
16
    txt_filename = os.path.join('yolo_seg_dataset/train/' + row[1].Path.split('/')[-1][:-4] + '.txt')
17
    with open(txt_filename, 'w') as up:
18
        for polygon in row[1].Polygons:
19
            normalized_polygon = normalize_polygon(polygon, img_width, img_height)
20
            normalized_coords = ' '.join([f'{coord[0]:.3f} {coord[1]:.3f}' for coord in normalized_polygon])
21
            up.write(f'0 {normalized_coords}\n')
22

23
 # 采用验证集
24
 for row in training_anno.iloc[10000:10150].iterrows():
25
    shutil.copy(row[1].Path, 'yolo_seg_dataset/valid')
26

27
    img = cv2.imread(row[1].Path)
28
    img_height, img_width = img.shape[:2]
29
    txt_filename = os.path.join('yolo_seg_dataset/valid/' + row[1].Path.split('/')[-1][:-4] + '.txt')
30
    with open(txt_filename, 'w') as up:
31
        for polygon in row[1].Polygons:
32
            normalized_polygon = normalize_polygon(polygon, img_width, img_height)
33
            normalized_coords = ' '.join([f'{coord[0]:.3f} {coord[1]:.3f}' for coord in normalized_polygon])
34
            up.write(f'0 {normalized_coords}\n')

逐行代码分析#

好的，我们逐行分析这段代码，它的主要功能是将训练集和验证集的数据从原始数据集中复制到新的目录中，并将多边形的坐标标准化为相对坐标。

1
if os.path.exists('yolo_seg_dataset'):
2
    shutil.rmtree('yolo_seg_dataset')

1
os.makedirs('yolo_seg_dataset/train')
2
os.makedirs('yolo_seg_dataset/valid')

1
def normalize_polygon(polygon, img_width, img_height):
2
    return [(x / img_width, y / img_height) for x, y in polygon]

1
# 采样训练集
2
for row in training_anno.iloc[:10000].iterrows():

1
    shutil.copy(row[1].Path, 'yolo_seg_dataset/train')

1
    img = cv2.imread(row[1].Path)
2
    img_height, img_width = img.shape[:2]

1
    txt_filename = os.path.join('yolo_seg_dataset/train/' + row[1].Path.split('/')[-1][:-4] + '.txt')

1
    with open(txt_filename, 'w') as up:

1
        for polygon in row[1].Polygons:
2
            normalized_polygon = normalize_polygon(polygon, img_width, img_height)

1
            normalized_coords = ' '.join([f'{coord[0]:.3f} {coord[1]:.3f}' for coord in normalized_polygon])

1
            up.write(f'0 {normalized_coords}\n')

1
# 采用验证集
2
for row in training_anno.iloc[10000:10150].iterrows():
3
    shutil.copy(row[1].Path, 'yolo_seg_dataset/valid')
4

5
    img = cv2.imread(row[1].Path)
6
    img_height, img_width = img.shape[:2]
7
    txt_filename = os.path.join('yolo_seg_dataset/valid/' + row[1].Path.split('/')[-1][:-4] + '.txt')
8
    with open(txt_filename, 'w') as up:
9
        for polygon in row[1].Polygons:
10
            normalized_polygon = normalize_polygon(polygon, img_width, img_height)
11
            normalized_coords = ' '.join([f'{coord[0]:.3f} {coord[1]:.3f}' for coord in normalized_polygon])
12
            up.write(f'0 {normalized_coords}\n')

步骤二：写入YOLO配置文件#

1
with open('yolo_seg_dataset/data.yaml', 'w') as up:
2
    data_root = os.path.abspath('yolo_seg_dataset/')
3
    up.write(f'''
4
path: {data_root}
5
train: train
6
val: valid
7

8
names:
9
    0: alter
10
''')

逐行代码分析#

这一段代码的目的是创建一个名为 data.yaml 的配置文件，通常用于YOLO等深度学习模型的数据配置。我们来逐步分析这段代码的结构和含义。

1
with open('yolo_seg_dataset/data.yaml', 'w') as up:

1
data_root = os.path.abspath('yolo_seg_dataset/')

1
up.write(f'''
2
path: {data_root}
3
train: train
4
val: valid
5

6
names:
7
    0: alter
8
''')

1
path: /home/user/projects/yolo_seg_dataset
2
train: train
3
val: valid
4

5
names:
6
    0: alter

步骤三：训练YOLO分割模型#

1
from ultralytics import YOLO
2

3
model = YOLO("./yolov8n-seg.pt")
4
results = model.train(data="./yolo_seg_dataset/data.yaml", epochs=10, imgsz=640)

逐行代码分析#

以下是对代码的逐行分析：

1
from ultralytics import YOLO

• 这行代码导入了 ultralytics 库中的 YOLO 类。
• ultralytics 是一个用于目标检测的库，提供了易于使用的接口来训练和推理 YOLO 模型，尤其是 YOLOv8 模型。
• YOLO 类是用于创建并管理 YOLO 模型的类，包括加载预训练模型、训练模型和进行推理。

1
model = YOLO("./yolov8n-seg.pt")

• 这行代码使用 YOLO 类创建了一个 model 实例。
• YOLO("./yolov8n-seg.pt") 加载一个预训练的 YOLOv8 模型，该模型文件是 yolov8n-seg.pt。这是 YOLOv8 模型的小型版本（n 代表 nano纳米），同时具有分割功能（-seg 代表分割任务）。
• 通过这种加载方式，可以直接基于预训练模型进行进一步训练或推理。

1
results = model.train(data="./yolo_seg_dataset/data.yaml", epochs=10, imgsz=640)

• model.train() 是 YOLO 模型的训练方法。
• data="./yolo_seg_dataset/data.yaml" 指定了用于训练的数据集配置文件。data.yaml 文件通常包含训练集、验证集的路径，以及类别名称等信息。
• epochs=10 指定模型要进行 10 个==训练周期==（epochs）。每个周期意味着模型将遍历完整的训练数据集一次。
• imgsz=640 设置了输入图像的尺寸为 640 像素。YOLO 模型会将训练数据的图像调整到 640x640 的尺寸进行训练。

训练完成后，results 会包含训练过程中的结果和性能指标（如损失值、准确度等），这些结果可以用于评估模型的训练效果。

步骤四：预测测试集#

1
from ultralytics import YOLO
2
import glob
3
from tqdm import tqdm
4

5
model = YOLO("./runs/segment/train6/weights/best.pt")
6
test_imgs = glob.glob('./test_set_A_rename/*/*')

1
Polygon = []
2
for path in tqdm(test_imgs[:]):
3
    results = model(path, verbose=False)
4
    result = results[0]
5
    if result.masks is None:
6
        Polygon.append([])
7
    else:
8
        Polygon.append([mask.astype(int).tolist() for mask in result.masks.xy])
9

10
import pandas as pd
11
submit = pd.DataFrame({
12
    'Path': [x.split('/')[-1] for x in test_imgs[:]],
13
    'Polygon': Polygon
14
})
15
submit.to_csv('track2_submit.csv', index=None)

逐行分析代码#

这段代码使用预训练的YOLO模型来处理图像集，并将结果保存到CSV文件中。下面是逐行的分析：

1
from ultralytics import YOLO
2
import glob
3
from tqdm import tqdm

• from ultralytics import YOLO: 从 ultralytics 库中导入 YOLO 类，这个类用于加载并使用预训练的 YOLO 模型进行推理和处理图像。
• import glob: glob 模块用于文件操作，它可以根据特定的模式匹配文件路径。
• from tqdm import tqdm: tqdm 模块用于显示进度条，常用于长时间运行的循环或任务中，帮助可视化处理进度。

1
model = YOLO("./runs/segment/train6/weights/best.pt")

• 这行代码加载了一个训练好的YOLO模型，权重文件位于 ./runs/segment/train6/weights/best.pt。这个模型是预训练的权重，专门用于图像分割任务。

1
test_imgs = glob.glob('./test_set_A_rename/*/*')

• 使用 glob.glob() 函数获取所有位于 ./test_set_A_rename/ 路径下的图像文件。
• '*/*' 匹配所有子目录中的文件路径，因此它会抓取该目录下所有的图像路径，并存储在 test_imgs 列表中。

1
Polygon = []
2
for path in tqdm(test_imgs[:]):
3
    results = model(path, verbose=False)
4
    result = results[0]

• Polygon = []: 初始化一个空列表 Polygon，用于存储每张图片的分割结果。
• for path in tqdm(test_imgs[:]):: 通过 tqdm 包裹循环，遍历 test_imgs 中的每个图像路径，并显示进度条。
• results = model(path, verbose=False): 调用模型对当前图片（path）进行推理，得到分割结果。verbose=False 表示不打印详细信息。
• result = results[0]: YOLO 返回的结果可能包含多张图片（批处理），这里仅取第一个结果（当前处理的图片）。

1
if result.masks is None:
2
    Polygon.append([])
3
else:
4
    Polygon.append([mask.astype(int).tolist() for mask in result.masks.xy])

• if result.masks is None:: 检查结果是否包含分割的掩码（masks）。如果没有分割掩码（即 result.masks 为 None），则向 Polygon 列表中添加一个空列表。

• Polygon.append([mask.astype(int).tolist() for mask in result.masks.xy]): 如果存在分割掩码，将掩码的边界坐标（result.masks.xy）转换为整数并存储为 Python 列表的形式，然后将其添加到 Polygon 列表中。

  ==掩码（Mask）==在计算机视觉中是一个非常重要的概念，它通常用于表示图像中的特定区域。掩码的作用是突出或分离出图像中的某些部分（例如物体、背景等），便于后续的处理。掩码通常是一个与原始图像尺寸相同的二值化矩阵，每个像素点要么是 1（代表属于目标区域），要么是 0（代表不属于目标区域）。

  在你的代码中，掩码（mask）与图像分割任务相关。模型返回的 mask 描述的是在图像中哪一部分属于分割出的对象。这些掩码可以是二值化的矩阵，也可以是其对应的轮廓坐标（多边形）。

1
import pandas as pd
2
submit = pd.DataFrame({
3
    'Path': [x.split('/')[-1] for x in test_imgs[:]],
4
    'Polygon': Polygon
5
})
6
submit.to_csv('track2_submit.csv', index=None)

• import pandas as pd: 导入 pandas 库，用于处理数据并生成 CSV 文件。
• submit = pd.DataFrame({...}): 创建一个 pandas 数据框 submit，包含两列：
  • 'Path': 存储每张图像的文件名（通过 x.split('/')[-1] 从路径中提取文件名）。
  • 'Polygon': 存储每张图像的分割结果（存储在 Polygon 列表中）。
• submit.to_csv('track2_submit.csv', index=None): 将 submit 数据框保存为一个名为 track2_submit.csv 的 CSV 文件，index=None 表示不保存行索引。