基于yolov5的车辆目标检测（二）数据集的获取

数据集的获取

一、将视频转换成图片

import os
import cv2


def video2img(video_folder, img_folder, frequency=10):
    video_files = os.listdir(video_folder)
    for video_file in video_files:
        print("{}图片提取中...".format(video_file))
        times = 0
        img_idx = 0
        video_file_name = video_file.split(".")[0]
        video_file_path = os.path.join(video_folder, video_file)
        img_sub_folder = os.path.join(img_folder, video_file_name)
        if not os.path.exists(img_sub_folder):
            os.makedirs(img_sub_folder)
            camera = cv2.VideoCapture(video_file_path)
            while True:
                times += 1
                res, image = camera.read()
                if not res:
                    print('EOF')
                    break
                if times % frequency == 0:
                    img_name = video_file_name + "_" + str(img_idx) + ".jpg"
                    img_idx = img_idx + 1
                    img_save_path = os.path.join(img_sub_folder, img_name)
                    cv2.imwrite(img_save_path, image)
            print('{}文件图片提取结束：{}张'.format(video_file, img_idx))
            camera.release()


if __name__ == '__main__':
    video2img(video_folder=r'E:\Machine-Vision\Test_Mov\Mov\MVI_2052', img_folder="E:\Machine-Vision\yolov5-master\data\images", frequency=2)

视频存放的路径为绝对路径，由于’E:\Machine-Vision\Test_Mov\Mov\MVI_2052’含有转义字符‘\'，所以这里在绝对路径前加r以区分转义字符
video_folder=视频的绝对路径
img_folder=照片保存的绝对路径
frequency=采样频率，相当于每隔frequency，对视频截取一次照片

这里选择的frequency=10

视频转换后的照片，保存在images路径下

二、对转换后的照片进行打标签

2.1使用lambeling软件进行打标签

这里将格式选择为yolo格式

2.2使用make sense进行打标签

make sense

点击Get Started


创建标签类别，选择为橙色，名称为car

将图片全部打完标签之后，点击上方菜单栏，Actions->Export Annotation->A.zip package containing files in YOLO format

之后生成zip压缩包，解压后为txt文件
将该文件重命名为labels，存放在与images相同的目录下

三、数据集的分类

数据集的分类使用到了如下代码：

import os
import shutil
import random

random.seed(0)


def split_data(file_path, xml_path, new_file_path, train_rate, val_rate, test_rate):
    each_class_image = []
    each_class_label = []
    for image in os.listdir(file_path):
        each_class_image.append(image)
    for label in os.listdir(xml_path):
        each_class_label.append(label)
    data = list(zip(each_class_image, each_class_label))
    total = len(each_class_image)
    random.shuffle(data)
    each_class_image, each_class_label = zip(*data)
    train_images = each_class_image[0:int(train_rate * total)]
    val_images = each_class_image[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_images = each_class_image[int((train_rate + val_rate) * total):]
    train_labels = each_class_label[0:int(train_rate * total)]
    val_labels = each_class_label[int(train_rate * total):int((train_rate + val_rate) * total)]
    test_labels = each_class_label[int((train_rate + val_rate) * total):]

    for image in train_images:
        print(image)
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'train' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in train_labels:
        print(label)
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'train' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

    for image in val_images:
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'val' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in val_labels:
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'val' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)

    for image in test_images:
        old_path = file_path + '/' + image
        new_path1 = new_file_path + '/' + 'test' + '/' + 'images'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + image
        shutil.copy(old_path, new_path)

    for label in test_labels:
        old_path = xml_path + '/' + label
        new_path1 = new_file_path + '/' + 'test' + '/' + 'labels'
        if not os.path.exists(new_path1):
            os.makedirs(new_path1)
        new_path = new_path1 + '/' + label
        shutil.copy(old_path, new_path)


if __name__ == '__main__':
    file_path = "E:\Machine-Vision\Data_get\data\car\images_car"
    xml_path = "E:\Machine-Vision\Data_get\data\car\labels_car"
    new_file_path = "E:\Machine-Vision\Data_get\data\car\imageSets_car"
    split_data(file_path, xml_path, new_file_path, train_rate=0.7, val_rate=0.2, test_rate=0.1)

其中：
images存放训练的图片
labels存放打好的标签.txt文件
imageSets存放分类好的数据集
file_path->images
xml_path->labels
new_file_path->imageSets

test 测试数据集
train 训练数据集
val 验证数据集