#split_data.py
#劃分數(shù)據(jù)集flower_data，數(shù)據(jù)集劃分到flower_datas中，訓練集：驗證集：測試集比例為6：2：2
import os
import random
from shutil import copy2
# 源文件路徑
file_path = r"D:/other/ClassicalModel/other/flower_data"
# 新文件路徑
new_file_path = r"D:/other/ClassicalModel/other/flower_datas"
# 劃分數(shù)據(jù)比例為6:2:2
split_rate = [0.6, 0.2, 0.2]
print("Starting...")
print("Ratio= {}:{}:{}".format(int(split_rate[0] * 10), int(split_rate[1] * 10), int(split_rate[2] * 10)))
class_names = os.listdir(file_path)
# 在目標目錄下創(chuàng)建文件夾
split_names = ['train', 'val', 'test']
# 判斷是否存在木匾文件夾
if os.path.isdir(new_file_path):
    pass
else:
    os.mkdir(new_file_path)
for split_name in split_names:
    # split_path = os.path.join(new_file_path, split_name)
    split_path = new_file_path + "/" + split_name
    if os.path.isdir(split_path):
        pass
    else:
        os.mkdir(split_path)
    # 然后在split_path的目錄下創(chuàng)建類別文件夾
    for class_name in class_names:
        class_split_path = os.path.join(split_path, class_name)
        if os.path.isdir(class_split_path):
            pass
        else:
            os.mkdir(class_split_path)
# 按照比例劃分數(shù)據(jù)集，并進行數(shù)據(jù)圖片的復制
# 首先進行分類遍歷
for class_name in class_names:
    current_class_data_path = os.path.join(file_path, class_name)
    current_all_data = os.listdir(current_class_data_path)
    current_data_length = len(current_all_data)
    current_data_index_list = list(range(current_data_length))
    random.shuffle(current_data_index_list)
    train_path = os.path.join(os.path.join(new_file_path, 'train'), class_name)
    val_path = os.path.join(os.path.join(new_file_path, 'val'), class_name)
    test_path = os.path.join(os.path.join(new_file_path, 'test'), class_name)
    train_stop_flag = current_data_length * split_rate[0]
    val_stop_flag = current_data_length * (split_rate[0] + split_rate[1])
    current_idx = 0
    train_num = 0
    val_num = 0
    test_num = 0
    for i in current_data_index_list:
        src_img_path = os.path.join(current_class_data_path, current_all_data[i])
        if current_idx <= train_stop_flag:
            copy2(src_img_path, train_path
            train_num = train_num + 1
        elif (current_idx > train_stop_flag) and (current_idx <= val_stop_flag):
            copy2(src_img_path, val_path)
            val_num = val_num + 1
        else:
            copy2(src_img_path, test_path
            test_num = test_num + 1
        current_idx = current_idx + 1
    print("<{}> has {} pictures,train:val:test={}:{}:{}".format(class_name, current_data_length, train_num, val_num,
                                                              test_num))
print("Done")

# *_*coding: utf-8 *_*
import os
import random
import shutil
import time
def copyFile(fileDir,origion_path1,class_name):
    name = class_name
    path = origion_path1
    image_list = os.listdir(fileDir) # 獲取圖片的原始路徑
    image_number = len(image_list)
    train_number = int(image_number * train_rate)
    train_sample = random.sample(image_list, train_number) # 從image_list中隨機獲取0.75比例的圖像.
    test_sample = list(set(image_list) - set(train_sample))
    sample = [train_sample, test_sample]
    # 復制圖像到目標文件夾
    for k in range(len(save_dir)):
            if os.path.isdir(save_dir[k]) and os.path.isdir(save_dir1[k]):
                for name in sample[k]:
                    name1 = name.split(".")[0] + '.xml'
                    shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k], name))
                    shutil.copy(os.path.join(path, name1), os.path.join(save_dir1[k], name1))
            else:
                os.makedirs(save_dir[k])
                os.makedirs(save_dir1[k])
                for name in sample[k]:
                    name1 = name.split(".")[0] + '.xml'
                    shutil.copy(os.path.join(fileDir, name), os.path.join(save_dir[k], name))
                    shutil.copy(os.path.join(path, name1), os.path.join(save_dir1[k], name1))
if __name__ == '__main__':
    time_start = time.time()
    # 原始數(shù)據(jù)集路徑
    origion_path = './JPEGImages/'
    origion_path1 = './Annotations/'
    # 保存路徑
    save_train_dir = './train/JPEGImages/'
    save_test_dir = './test/JPEGImages/'
    save_train_dir1 = './train/Annotations/'
    save_test_dir1 = './test/Annotations/'
    save_dir = [save_train_dir, save_test_dir]
    save_dir1 = [save_train_dir1, save_test_dir1]
    # 訓練集比例
    train_rate = 0.75
    # 數(shù)據(jù)集類別及數(shù)量
    file_list = os.listdir(origion_path)
    num_classes = len(file_list)
    for i in range(num_classes):
        class_name = file_list[i]
    copyFile(origion_path,origion_path1,class_name)
    print('劃分完畢!')
    time_end = time.time()
    print('---------------')
    print('訓練集和測試集劃分共耗時%s!' % (time_end - time_start))