目标检测划分数据集

目标检测划分数据集

训练检测在数据采集、数据标注后需要划分数据集，一般对标注的数据集划分为训练集和验证集，用于一些通用算法 的训练，如YOLO、SSD、faster-RCNN等

# -*- coding:utf8 -*-
import time
import os
import random
import cv2


t_1 = time.time()
imgs_path = './drone-roi-low/images/train/'
anns_path = './drone-roi-low/labels/train/'

name = 'fruits_dataset' # 划分后数据集文件夹名，随意更改
img_train_path ='./'+name+'/images/train/'  #图片训练集保存路径
img_val_path = './'+name+'/images/val/'     #图片验证集保存路径
labels_train_path = './'+name+'/labels/train/'       #标签训练集保存路径
labels_val_path = './'+name+'/labels/val/'           #标签验证集保存路径
if not os.path.exists(img_train_path):
    os.makedirs(img_train_path)
if not os.path.exists(img_val_path):
    os.makedirs(img_val_path)
if not os.path.exists(labels_train_path):
    os.makedirs(labels_train_path)
if not os.path.exists(labels_val_path):
    os.makedirs(labels_val_path)

imgs = os.listdir(imgs_path)
anns = os.listdir(anns_path)
random.seed(2021)                           #设置一个随机种子，确保每次运行都按照既定的随机形式
random.shuffle(imgs)
train_set_rate = 0.8
train_set_num = int(len(imgs) * train_set_rate)

train_name_list = []
val_name_list = []
num_img = 0
num_train_set = 0
num_val_set = 0
for i ,img in enumerate(imgs) :
    image = cv2.imread(imgs_path + img)
    if i <= train_set_num:
        img = img.split('.')
        train_name_list.append(img[0])
        img = str.join('.', img)
        cv2.imwrite(img_train_path + img , image)
        num_train_set+=1

    else:
        img = img.split('.')
        val_name_list.append(img[0])
        img = str.join('.',img)
        cv2.imwrite(img_val_path+ img, image)
        num_val_set+=1
    i+=1
    num_img+=1
    print('num_img----------------------',i)

print('train_set--------------',num_train_set)
print('val_set-----------------',num_val_set)
print('图片总数量-----------------------',num_img)

ann_train = 0
ann_val = 0
num = 0
for ann in anns:
    ann = ann.split('.')
    if (ann[0] in train_name_list) :
        ann = str.join('.',ann)
        with open(anns_path + ann ,'r', encoding = 'UTF-8')as f:
            with open(labels_train_path + ann ,'w')as s:
                s.write(f.read())
        ann_train+=1

    elif (ann[0] in val_name_list):
        ann = str.join('.', ann)
        with open(anns_path + ann, 'r',encoding = 'UTF-8')as f:
            with open(labels_val_path + ann, 'w')as s:
                s.write(f.read())
        ann_val+=1
    num+=1
t_2 = time.time()
print('ann_train--------------------',ann_train)
print('ann_val----------------------',ann_val)
print('totle_num---------------------',num)
print('time==========================',t_2 - t_1)