python 随机划分图片数据集以及移动标注

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')