YOLO数据集制作

1、视频提取图片

#! /usr/bin/python
# -*- coding:utf-8 -*-
import cv2
import os
import numpy as np
import random 
import shutil


def video_2_images(video_path, save_iamge_path):
    count = 0
    for v_path in video_path:
        print('is dealing {}'.format(v_path))
        video = cv2.VideoCapture(v_path)
        while video.isOpened():
            ret, frame = video.read()
            if not ret:
                break
            cv2.imwrite(os.path.join(save_iamge_path, str(count)+'.jpg'), frame)
            count += 1
        video.release()

def write_2_txt(txt_path, label_data_list):
    '''
    label_data_list : 列表， 元素是 str, 每个str是“cls x y w h” xywh是除以图像宽高之后的数值
    '''
    with open(txt_path, 'a') as f:
        f.writelines(label_data_list)
 
           
def use_cv2_label_image(images_path, save_label_path):
    all_images_name = os.listdir(images_path)
    for image_name in all_images_name:
        print('doing {}'.format(image_name))
        img_path = os.path.join(images_path, image_name)
        save_txt_path = os.path.join(save_label_path, image_name.split('.')[0] + '.txt')

        img = cv2.imread(img_path, 0)
        _, img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)  
        kernel = np.ones((5,5), np.uint8)
        img = cv2.dilate(img, kernel, iterations = 1)
        # img = cv2.erode(img, kernel, iterations = 1)
        
        
        contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        areas = []
        select_contours = []
        for c in contours:
            area = cv2.contourArea(c)
            if area > 3000 and area < 4500:
                areas.append(area)
                select_contours.append(c)
        if select_contours == []:
            continue
        else:
            left_up_x, left_up_y = 0., 0.
            right_down_x, right_down_y = 0., 0.
            w, h = 0., 0.
            img_w = img.shape[1]
            img_h = img.shape[0]
            label_rect = []
            
            # draw_img = np.zeros(img.shape, np.uint8)
            for c in select_contours:
                c = np.squeeze(c)
                left_up_x = np.min(c, axis=0)[0]
                left_up_y = np.min(c, axis=0)[1]
                right_down_x = np.max(c, axis=0)[0]
                right_down_y = np.max(c, axis=0)[1]
                w = right_down_x - left_up_x
                h = right_down_y - left_up_y
                label_rect.append('0 '+ str(round((left_up_x + right_down_x)/(2*img_w), 6)) + ' '+ \
                    str(round((left_up_y + right_down_y)/(2 * img_h), 6)) + ' '+ \
                        str(round(w/img_w, 6)) + ' ' + str(round(h/img_h, 6)) )
                
                # cv2.rectangle(draw_img, (left_up_x, left_up_y), (right_down_x, right_down_y), (255,255,255), -1)
            write_2_txt(save_txt_path, label_rect)
        # print("img {}, area = {}".format(image_name, areas))
        # draw_img = np.zeros(img.shape, np.uint8)
        # cv2.drawContours(draw_img, select_contours, -1, (255,255,255), thickness = -1)
        # cv2.imwrite(os.path.join("/home/nvidia/rk3588/dataset/label_images", image_name), draw_img)

def allocate_dataset(source_label_path, source_image_path, target_label_path, target_image_path, train_rate=0.6):
    all_label_names = os.listdir(source_label_path)
    all_labeled_iamges_path = [os.path.join(source_image_path, label_name.split('.')[0] + '.jpg') for label_name in all_label_names]
    all_labels_path = [os.path.join(source_label_path, label_name) for label_name in all_label_names]
    
    # 按照比例随即分配数量给train val
    all_labels_num = len(all_label_names)
    train_num = int(all_labels_num * train_rate)
    train_sample_index = random.sample(range(all_labels_num), train_num)
    
    for i in range(all_labels_num):
        if i  in train_sample_index:
            shutil.move(all_labeled_iamges_path[i], os.path.join(target_image_path, 'train'))
            shutil.move(all_labels_path[i], os.path.join(target_label_path, 'train'))
        else:
            shutil.move(all_labeled_iamges_path[i], os.path.join(target_image_path, 'val'))
            shutil.move(all_labels_path[i], os.path.join(target_label_path, 'val'))
         

if __name__ == '__main__':
    # video_2_images(video_path=["/home/nvidia/l/30.mp4", "/home/nvidia/l/20.mp4"], save_iamge_path="/home/nvidia/rk3588/dataset/original_images")
    # use_cv2_label_image(images_path="/home/nvidia/rk3588/dataset/original_images", save_label_path="/home/nvidia/rk3588/dataset/original_labels")
    allocate_dataset(source_label_path="/home/nvidia/rk3588/dataset/original_labels", source_image_path="/home/nvidia/rk3588/dataset/original_images", target_label_path="/home/nvidia/rk3588/dataset/tank/labels", target_image_path="/home/nvidia/rk3588/dataset/tank/images")

2、LabelImg制作标签

1、平台/环境

jetson xavier , aarch64 ， conda（pyhton=3.6环境）

2、安装 PyQt5 （pyhton=3.6环境）

(1)使用 pip3 install PyQt5 在 aarch平台上装不上，一直卡住；
(2) 正确方法，参考这里；
主要步骤：

sudo apt install pyqt5*
sudo apt-get install python3-pyqt5

建立软链接，不可直接复制，根据自己环境修改
ln -s /usr/lib/python3/dist-packages/PyQt5/  ~/miniforge3/envs/mytest/lib/python3.6/site-packages
ln -s /usr/lib/python3/dist-packages/sip*  ~/miniforge3/envs/mytest/lib/python3.6/site-packages

3、下载LabelImg源码

链接
使用：

python3 labelImg.py

LabelImg 使用方法在此。

4、XML文件转成 txt 文件

由于LabelImg 标注完之后生成的是 xml 格式的文件， yolo 训练需要的是txt 格式的文件，所以需要一个脚本转换，下面代码转载自他人，出处一时半会找不到了，见谅。

# -×- coding:utf-8 -*-
import os
from os import listdir,getcwd
import re

classes=[]					#列表，储存标签类别
classes_dict={}				#字典，储存标签类别及对应的序号，yolo标签文件中用序号表示各个类别

def convert(Imgsize,box):
    width,height=Imgsize
    xmin,ymin,xmax,ymax=box
    dw=max(xmin,xmax)-min(xmin,xmax)#框的宽
    dh=max(ymin,ymax)-min(ymin,ymax)#框的高
    x_=(xmin+xmax)/2.0#框的中心点x坐标
    y_=(ymin+ymax)/2.0#框的中心点y坐标

    #归一化
    cx=round(x_/width,6)
    cy=round(y_/height,6)
    w=round(dw/width,6)
    h=round(dh/height,6)

    return [cx,cy,w,h]


def save_txt(namepath,text):			#将内容保存至文本文件中，即成功转化为yolo的标签文件
    with open(namepath,'w')as f:
        f.write(text)

def convert_annotation(xml_path,name):
    xml_name=xml_path+"/"+name
    with open(xml_name,"r",encoding="utf-8")as f1:
        text=f1.read().replace("\n","")
        text=text.replace(" ","")
    # print(text)
    img_size=re.findall("<width>([0-9]+)</width>.*?<height>([0-9]+)</height>",text)[0]
    print(img_size)
    find_datas=re.findall("<object>.*?<name>([a-z|A-Z]*?)</name>.*?<xmin>([0-9]+?)</xmin>.*?<ymin>([0-9]+?)</ymin>.*?<xmax>([0-9]+?)</xmax>.*?<ymax>([0-9]+?)</ymax>",text)
    # print(find_datas)
    savetext=""
    for item in find_datas:			#由于每张图片中可能存在多个类别或标签，将同张图片中所有标签内容保存在一个txt文件中
        class_=item[0]				#类别名
        if class_ not in classes:	#判断此类别是否已存在于classes列表中
            classes.append(class_)	#如果不存在，则添加至列表中
            classes_dict[class_]=len(classes)-1		#同时，为此类别标记序号，即第几个类别，从0开始标号

        print(item)
        imgsize=[int(img_size[0]),int(img_size[1])]			#正则获取的内容为文本格式，需要转化为数字
        box=[int(item[1]),int(item[2]),int(item[3]),int(item[4])]
        site=convert(imgsize,box)							#整理内容，归一化数值
        # print(class_,site)
        savetext+="{0} {1} {2} {3} {4}".format(classes_dict[class_],site[0],site[1],site[2],site[3])		#按格式拼接内容
        savetext+="\n"			#文本换行
    # print(savetext)
    name=name.split(".")[0]		#获取文件名，不包含文件后缀
    save_txt(labels_p+"/"+name+".txt",savetext.strip())			#保存标签文件
    

if __name__=="__main__":
    root_path=os.getcwd()			#获取当前目录绝对路径
    print(root_path)
    labels_p= '/home/nvidia/Dataset/txt_label'	#定义标签储存的文件夹路径
    
    xml_path= '/home/nvidia/Dataset/xml_label'		#.xml文件的父级文件夹路径
    xml_list=sorted(listdir(xml_path))			#获取所有.xml文件名，并保存至列表中，同时排序
    print(xml_list)
    
    for name in xml_list:						#历遍列表
        print(name)
        convert_annotation(xml_path,name)		#开始转化
    
    print(classes,classes_dict)
    with open(labels_p+"/classes.txt","w")as f2:	#保存类别标签文件
        text=""
        for t in classes:
            text+=t
            text+="\n"
        print(text)
        f2.write(text)