首页 > 解决方案 > 为 YOLOv3 分割基于图像的数据集

问题描述

我有一个关于拆分 20k 图像及其标签的数据集的问题,该数据集的格式YOLOv3具有一个图像文件和一个.txt与图像同名的文件,文本文件里面有标签。

我想将数据集拆分为训练/测试拆分,有没有办法随机选择图像及其标签 .txt 文件并使用它存储在单独的文件夹中Python

我希望能够随机拆分数据集。例如,选择 16k 文件和标签文件,并将它们单独存储在 train 文件夹中,剩余的 4k 应存储在 test 文件夹中。

这可以在文件资源管理器中手动完成,方法是选择前 16k 个文件并将它们移动到不同的文件夹,但拆分不会是随机的,因为我计划对同一个数据集一遍又一遍地执行此操作。

这是数据看起来像 图像和标签的截图

标签: pythondatasetconv-neural-networkyolo

解决方案


我建议你看看下面的 Python 内置模块

用于在 Python 中操作文件和路径。这是我的代码,其中包含可能解决您问题的注释。这很简单

import glob
import random
import os
import shutil

# Get all paths to your images files and text files
PATH = 'path/to/dataset/'
img_paths = glob.glob(PATH+'*.jpg')
txt_paths = glob.glob(PATH+'*.txt')

# Calculate number of files for training, validation
data_size = len(img_paths)
r = 0.8
train_size = int(data_size * 0.8)

# Shuffle two list
img_txt = list(zip(img_paths, txt_paths))
random.seed(43)
random.shuffle(img_txt)
img_paths, txt_paths = zip(*img_txt)

# Now split them
train_img_paths = img_paths[:train_size]
train_txt_paths = txt_paths[:train_size]

valid_img_paths = img_paths[train_size:]
valid_txt_paths = txt_paths[train_size:]

# Move them to train, valid folders
train_folder = PATH+'train/' 
valid_folder = PATH+'valid/'
os.mkdir(train_folder)
os.mkdir(valid_folder)

def move(paths, folder):
    for p in paths:
        shutil.move(p, folder)

move(train_img_paths, train_folder)
move(train_txt_paths, train_folder)
move(valid_img_paths, valid_folder)
move(valid_txt_paths, valid_folder)


推荐阅读