現有一個零售商品圖片的數據集goods,沒有做任何劃分。goods目錄下,子目錄的名稱為圖片的類別,每個類別下有多張圖片,文件結構如圖1。現在,我們需要按一定的比例把它拆分為訓練集和測試集。
實現該功能的思路是建立一個新的相同結構的目錄樹,並在原目錄樹中的每個類別中隨機按比例地挑選部分圖片,將它們移動到新目錄樹中相應的類別下。完整代碼如下:
import os
import random
import shutil
#源數據集路徑和目標數據集路徑
path_source = './goods'
path_target = './goods_test'
#參數:源路徑、目標路徑和測試集所占比例
def seperate(path_source, path_target, percent):
#生成包含path_source下所有目錄名的列表
categories = os.listdir(path_source)
for name in categories:
#在path_target下建立相同名稱的子目錄
os.makedirs(os.path.join(path_target, name))
#生成包含子目錄下所有圖片的列表
nums = os.listdir(os.path.join(path_source, name))
#隨機按比例抽取一部分圖片
nums_target = random.sample(nums, int(len(nums)*percent))
#把圖片剪切到目標路徑
for pic in nums_target:
shutil.move(os.path.join(path_source, name, pic), os.path.join(path_target, name, pic))
#執行完成後,path_source為訓練集,path_target為測試集。
seperate(path_source, path_target, 0.3)
如果數據集是其他結構,在此基礎上稍作改動即可。
其中的關鍵函數:
os.listdir(path):列出path下所有項目的名稱(包括目錄和文件)
os.makedirs(path):遞歸地創建path路徑上的目錄,目錄已經存在的話會報錯
os.makedir(path)可以用於僅創建一個目錄的情況
random.sample(list, num):從list中隨機抽取num個元素組成新的列表
shutil.move(source, target):將source文件移動到target文件
注意shutil.move是個很危險的操作,建議在調試代碼時注釋掉,確保代碼無誤後再執行。