Python实现随机划分图片数据集的示例代码

一、前言

前几天需要划分数据集，就写了一个小demo去完成这个任务。

随机划分图片数据集

任务描述：我的所有图片保存在同一个文件夹里，需要随机将图片与标注文件划分为训练集和测试集。

处理过程：读取文件列表，将列表打乱，截取列表一部分

二、实现代码如下

import os
import random
import shutil

def get_imlist(path):
   return [os.path.join(path, f) for f in os.listdir(path) if f.endswith('.jpg')]

def getData(src_path):

dest_dir = src_path+'val' #划分出来的验证集
   if not os.path.isdir(dest_dir):
       os.mkdir(dest_dir)

img_list = get_imlist(src_path)
   random.shuffle(img_list)
   le = int(len(img_list) * 0.8)  # 这个可以修改划分比例
   for f in img_list[le:]:
       shutil.move(f, dest_dir)

'''
函数功能：
划分数据集
'''
def SplitImg(filePath):
   getData(filePath)

'''
函数功能：
根据划分的数据集进行移动标注文件
'''
def MoveAn(filePathAn,filePathImg):
   Imgs=os.listdir(filePathImg)
   if not os.path.isdir(filePathAn+'val'):
       os.mkdir(filePathAn+'val')

for file in os.listdir(filePathAn):
       #print(filePathAn,filePathImg)
       #print(os.path.join(filePathAn,file),os.path.join(filePathAn+'val',file))
       if file[:-4]+'.jpg' in Imgs:

shutil.move(os.path.join(filePathAn,file),os.path.join(filePathAn+'val',file))

if __name__=='__main__':
   filePath='./宠物/宠物'# 换成你的数据集

#拆分的数据集
   SplitImg(filePath)
   filePathAn='./宠物/宠物An'# 换成你的标注文件地址

# 根据数据集进行移动标注文件
   MoveAn(filePathAn,filePath+'val')

三、方法补充

除了以上的方法，小编还为大家整理了其他划分数据集的方法，希望对大家有所帮助

方法一：使用random.sample将数据集随机划分为训练集与验证集并另存在文件夹中，设置随机种子

import os
import random
import shutil

def moveFile(input1,input2,save1,save2):
   pathDir = os.listdir(input1)  # 取图片的原始路径
   random.seed(1)
   filenumber = len(pathDir)  # 原文件个数
   rate = 0.1  # 抽取的验证集的比例，占总数据的多少
   picknumber = int(filenumber * rate)  # 按照rate比例从文件夹中取一定数量图片
   sample = random.sample(pathDir, picknumber)  # 随机选取需要数量的样本图片
   print(sample)
   list_len=len(sample)
   print(list_len)
   list=[]
   for i in range(len(sample)):
       list.append(sample[i].split('.')[0])
   print(list)
   for flie_name in list:
       path_img=os.path.join(input1,flie_name+'.jpg')
       shutil.move(path_img,save1)
       path_lab=os.path.join(input2,flie_name+'.txt')
       shutil.move(path_lab,save2)

if __name__ == '__main__':
   input_path1='./train1200/images/train'
   input_path2= './train1200/labels/train'
   save_img='./train1200/images/val'
   save_lab='./train1200/labels/val'
   if not os.path.exists(save_lab):
       os.makedirs(save_lab)
   if not os.path.exists(save_img):
       os.makedirs(save_img)
   moveFile(input_path1,input_path2,save_img,save_lab)

方法二：通过train test_splt函数实现随机划分数据

Python中，随机划分数据集可以通过train test_splt函数实现。该函数可以将数据集随机分成训练集和测试集，用于机器学习中的数据训练和测试。

函数的基本用法如下:

from sklearn.model_selection import train_test_split

#X是所有的样本特征，y是目标变量，test_size是测试集占总样本数的比例
# random_state是随机数发生器的种子，保证每次划分结果一致
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0,3, random_state=42)

其中，X是所有的样本特征，是目标变量，test_size是测试集占总样本数的比例。random_state用于设置随机数发生器的种子，以保i每次划分结果一致。函数会返回4个数组: X_train、X_test、y_train,y_test，分别代表训练集的特征、测试集的特征、训练集的目标变量和测试集的目标变量

方法三：将一个数据集按比例随机分割成训练集、验证集、测试集

import random

def split(fname, train_ratio, var_ratio):
   lines = fname.readlines()
   n_total = len(lines)          # 获取数据集的总长度

train_offset = int(n_total * train_ratio)
   val_offset = int(n_total * (train_ratio + var_ratio))
   random.shuffle(fname.read())  # 按行打乱顺序

train_data = open('train.txt.bio', 'wb')
   val_data = open('val.txt.bio', 'wb')
   test_data = open('test.txt.bio', 'wb')

for i, line in enumerate(lines):
       if i < train_offset:
           train_data.write(line)
       elif i < val_offset:
           val_data.write(line)
       else:
           test_data.write(line)

train_data.close()
   val_data.close()
   test_data.close()

if __name__ == "__main__":
  fname = open('en/en_total.txt.bio', "rb")
  split(fname, train_ratio = 0.6, var_ratio = 0.2)
  fname.close()

来源：https://blog.csdn.net/qq_51116518/article/details/126002379