如何快速为YOLO准备标注好的图像样本库

1. YOLO目标检测

    YOLO是时下工业界用的比较多的视觉目标检测算法，特别是YOLOv5有专门的团队维护，提供了从小s到大l到巨x各种预训练模型文件，并且为使用者提供了近乎傻瓜式的train和detect，是非常良心的github项目。不过，YOLOv5要求提供满足格式要求的标注图像样本，具体而言：

1. 要求提供yaml文件，来描述样本图片的路径，样本的类别数量，以及每种类别的含义；

2. 样本图片和每张图片同名txt标准文件，文件的每行为（详见官方文档中的 Or manually prepare your dataset (click to expand) 记得点开）：

class x_center y_center width heigh

    关于YOLOv5如何跑自己的数据集，B站up主做了详细讲解。只不过，在图片上面一个一个去框选标注有些费力，或许可以试试通过视频文件自动生成标注图片？

2. 通过视频生成标注样本库的思路

    假设我们的被检测对象是一个物体，我们对这类物体录制了多段视频，那如果我们能通过先验知识+图像处理手段从视频中找出该物体的x,y,w,h，不就能生成标注文件了么？30帧的视频1分钟就是1800张图片啊，并且还有丰富的动态特性！下面就是如何找出视频中物体的边框了，这就要通过opencv的findContours了。

3. opencv找出物体轮廓

    findContours是opencv十分强大的函数，普通图片经过高斯滤噪、色彩过滤等预处理后，就可以找到轮廓了。一次会得到不止一个轮廓，这时候就要预先给定物体出现的区域中心点位置，然后把包含这个中心点的轮廓作为样本轮廓，然后调用cv2.contourArea(contour)得到轮廓的x,y,w,h，简单吧。

4.用x,y,w,h生成YOLOv5标注

    前面提的B站up主提供了一段代码来做转换convert，无非就是除以长宽得到相对坐标，再用x=x+w/2，y=y+h/2 得到中心点坐标，最后合并得到 class centerx centery width height 这样一条标注信息，写入到txt文件。

5. 例子代码

    下面这段代码在dataset\flame2\images下生成图片，在dataset\flame2\labels下生成对应的标注文件。由于是单类对象检测问题，class只有一个，所以标注只有一行，且class=0。

import cv2
import numpy as np


def convert(dw, dh, x, y, w, h):
    x += w/2.0
    y += h/2.0
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return f'0 {x} {y} {w} {h}'


videoname = 'flame2'
cap = cv2.VideoCapture(videoname+'.mp4')
success, frame = cap.read()
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
dh = 1.0/height
dw = 1.0/width

centerx, centery = 585, 231
lower_bound = np.array([0, 0, 175])
upper_bound = np.array([360, 120, 249])
imagepath = f'dataset\\{videoname}\\images\\'
labelpath = f'dataset\\{videoname}\\labels\\'
frameindex = 0
stat = {}
while success:
    blur = cv2.GaussianBlur(frame, (3, 3), 0)
    frame_hsv = cv2.cvtColor(blur, cv2.COLOR_BGR2HSV)
    frame_copy = cv2.inRange(frame_hsv, lower_bound, upper_bound)
    contours, hier = cv2.findContours(frame_copy, cv2.RETR_EXTERNAL,
                                      cv2.CHAIN_APPROX_SIMPLE)

    if len(contours) > 0:
        maxcontour = None
        maxarea = 0
        mx, my, mw, mh = 0, 0, 0, 0
        for i, c in enumerate(contours):
            x, y, w, h = cv2.boundingRect(c)
            if centerx > x and centerx < x+w and centery > y and centery < y + h:
                area = cv2.contourArea(c)
                if maxarea < area:
                    maxarea = area
                    maxcontour = c
                    mx, my, mw, mh = x, y, w, h
        if maxarea > 100 and maxcontour is not None:
            imagename = f'{imagepath}{videoname}_{frameindex}.jpg'
            textname = f'{labelpath}{videoname}_{frameindex}.txt'
            cv2.imwrite(imagename, frame)
            file2 = open(textname, 'w+')
            file2.write(convert(dw, dh, mx, my, mw, mh))
            file2.close()
            frameindex += 1
    cv2.imshow("Fire", frame)
    if cv2.waitKey(10) == 27:
        break
    success, frame = cap.read()
cap.release()
cv2.destroyAllWindows()

6. 使用效果

    我用手机录的几段视频，几分钟就生成了上万张标注图像，然后YOLOv5 train 起来！颇有些一夜暴富的感觉，你还等什么？