Python實現數據集劃分(訓練集和測試集)

更新時間：2023年05月19日 17:06:07 作者：Python小丸子

這篇文章主要為大家詳細介紹了Python是如何實現數據集劃分的，分為訓練集和測試集，文中的實現方法講解詳細，感興趣的小伙伴可以了解一下

前面是分部講解，完整代碼在最后。

導入模塊：

import os
from shutil import copy, rmtree
import random

創(chuàng)建文件夾：

def make_file(file_path: str):
    if os.path.exists(file_path):
        rmtree(file_path)
    os.makedirs(file_path)

劃分數據集的比例，本文是0.1：驗證集的數量占總數據集的10%比如填0.1就是驗證集的數量占總數據集的10%

random.seed(0)  
split_rate = 0.1

數據集的存放：新建一個數據文件夾，將劃分的數據集存放進去

data_path = r'D:\chengxu\data\caodi'  # 數據集存放的地方
data_root = r'D:\chengxu\data\cd'  # 這里是生成的訓練集和驗證集所處的位置，這里設置的是在當前文件夾下。
data_class = [cla for cla in os.listdir(data_path)]
print("數據的種類分別為：")
print(data_class)  # 輸出數據種類

建立訓練集文件夾：

train_data_root = os.path.join(data_root, "train")  # 訓練集的文件夾名稱為 train
make_file(train_data_root)
for num_class in data_class:
    make_file(os.path.join(train_data_root, num_class))

建立測試集文件夾：

val_data_root = os.path.join(data_root, "val")  # 驗證集的文件夾名稱為 val
make_file(val_data_root)
for num_class in data_class:
    make_file(os.path.join(val_data_root, num_class))

劃分數據：

for num_class in data_class:
    num_class_path = os.path.join(data_path, num_class)
    images = os.listdir(num_class_path)
    num = len(images)
    val_index = random.sample(images, k=int(num * split_rate))  # 隨機抽取圖片
    for index, image in enumerate(images):
        if image in val_index:
            # 將劃分到驗證集中的文件復制到相應目錄
            data_image_path = os.path.join(num_class_path, image)
            val_new_path = os.path.join(val_data_root, num_class)
            copy(data_image_path, val_new_path)
        else:
            # 將劃分到訓練集中的文件復制到相應目錄
            data_image_path = os.path.join(num_class_path, image)
            train_new_path = os.path.join(train_data_root, num_class)
            copy(data_image_path, train_new_path)
    print("\r[{}] split_rating [{}/{}]".format(num_class, index + 1, num), end="")  # processing bar
    print()
print("       ")
print("       ")
print("劃分成功")

完整代碼：

import os
from shutil import copy, rmtree
import random
 
 
def make_file(file_path: str):
    if os.path.exists(file_path):
 
        rmtree(file_path)
    os.makedirs(file_path)
 
random.seed(0) 
 
# 將數據集中10%的數據劃分到驗證集中
split_rate = 0.1  
data_path = r'D:\chengxu\data\caodi'  # 數據集存放的地方，建議在程序所在的文件夾下新建一個data文件夾，將需要劃分的數據集存放進去
data_root = r'D:\chengxu\data\cd'  # 這里是生成的訓練集和驗證集所處的位置，這里設置的是在當前文件夾下。
 
data_class = [cla for cla in os.listdir(data_path)]
print("數據的種類分別為：")
print(data_class)  
# 建立保存訓練集的文件夾
train_data_root = os.path.join(data_root, "train")  # 訓練集的文件夾名稱為 train
make_file(train_data_root)
for num_class in data_class:
    # 建立每個類別對應的文件夾
    make_file(os.path.join(train_data_root, num_class))
 
# 建立保存驗證集的文件夾
val_data_root = os.path.join(data_root, "val")  # 驗證集的文件夾名稱為 val
make_file(val_data_root)
for num_class in data_class:
    # 建立每個類別對應的文件夾
    make_file(os.path.join(val_data_root, num_class))
 
for num_class in data_class:
    num_class_path = os.path.join(data_path, num_class)
    images = os.listdir(num_class_path)
    num = len(images)
 
    val_index = random.sample(images, k=int(num * split_rate))  # 隨機抽取圖片
    for index, image in enumerate(images):
        if image in val_index:
          
            data_image_path = os.path.join(num_class_path, image)
            val_new_path = os.path.join(val_data_root, num_class)
            copy(data_image_path, val_new_path)
        else:
      
            data_image_path = os.path.join(num_class_path, image)
            train_new_path = os.path.join(train_data_root, num_class)
            copy(data_image_path, train_new_path)
    print("\r[{}] split_rating [{}/{}]".format(num_class, index + 1, num), end="")  # processing bar
    print()
 
print("       ")
print("       ")
print("劃分成功")

到此這篇關于Python實現數據集劃分(訓練集和測試集)的文章就介紹到這了,更多相關Python數據集劃分內容請搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關文章希望大家以后多多支持腳本之家！

您可能感興趣的文章: