目录
1. YOLO简介
2.分割数据集准备
3.代码展示
整理不易,欢迎一键三连!!!
1. YOLO简介
YOLO(You Only Look Once)是一种流行的目标检测和图像分割模型,由华盛顿大学的 Joseph Redmon 和 Ali Farhadi 开发。YOLO 的第一个版本于 2015 年发布,并因其高速度和准确性而迅速流行起来。
YOLO不同版本发行时间
版本 | 时间 |
YOLOv1 | 2015 |
YOLOv2 | 2016 |
YOLOv3 | 2018 |
YOLOv4 | 2020 |
YOLOv5 | 2021 |
YOLOv8 | 2022 |
以YOLOv5为例,说明同时支持图像分类、目标识别、图像分割的多任务网络训练的数据准备工作。这几天在网上找了好几天,一边摸索,一边准备自己的数据集,终于搞定了,记录下保姆级数据集准备的教程。
2.分割数据集准备
通常的分割任务中数据集都是一张原始影像对应一张同样尺寸的掩模图像,YOLO最开始是跑目标检测任务,因此常用的数据组织格式都是一张原始影像对应一个json掩膜文件或者一个TXT掩膜文件,这个看不同的目标检测数据格式就知道(COCO /VOC/......),今天以一张影像对应一个txt掩膜为例来做自己的数据集。
------>
从原始的png掩膜转为YOLO需要的txt掩膜文件关键是找到txt文件中的内容和组织方式:
从上面实例txt文件 可以看到, 第一行的第一个值为“45”,代表类别是45,后面的一堆小数点(0.78...)代表被归一化后的多边形的x,y坐标,归一化的标准是按照原始影像大小进行归一化,比如原始像素坐标为(10,20),原始影像尺寸为100*100,那归一化后的像素坐标为(0.1,0.2)。
回车之后,第二行开始读取第二个目标的类别和位置,以此类推。
3.代码展示
from skimage import io
import cv2
import numpy as np
def mask_to_polygon(mask: np.array, report: bool = False) -> List[int]:
contours, _ = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
polygons = []
for object in contours:
coords = []
for point in object:
coords.append(int(point[0][0]))
coords.append(int(point[0][1]))
polygons.append(coords)
if report:
print(f"Number of points = {len(polygons[0])}")
return np.array(polygons).ravel().tolist()
mask = io.imread('/labels/xxx.png')
polygons = mask_to_polygon(mask, report=True)
函数返回得到的polygons就是一个存放了一张影像中所有目标的多边形位置了,要想得到最终的txt文件,别忘了在前面加上影像类别哦。
如果像简化掩膜要素,就是把距离较近的点删除一些,减少掩膜文件的大小,可以参考下面的文章。
参考:二进制掩膜转txt
官方给出的COCO128-seg数据集:下载链接