Untitled
import os import shutil import random def split_data(source_dir, train_dir, test_dir, split_ratio=0.8): """ Chia dữ liệu từ source_dir thành train_dir và test_dir theo tỷ lệ split_ratio. Bên trong mỗi folder train và test sẽ có cùng cấu trúc thư mục với source_dir. :param source_dir: Thư mục nguồn chứa các thư mục nhỏ cần chia. :param train_dir: Thư mục đích để lưu tập huấn luyện. :param test_dir: Thư mục đích để lưu tập kiểm thử. :param split_ratio: Tỷ lệ chia tập huấn luyện và kiểm thử (mặc định là 0.8). """ # Lặp qua tất cả các thư mục con trong source_dir for folder in os.listdir(source_dir): folder_path = os.path.join(source_dir, folder) if os.path.isdir(folder_path): # Tạo các thư mục tương ứng trong train và test train_folder = os.path.join(train_dir, folder) test_folder = os.path.join(test_dir, folder) os.makedirs(train_folder, exist_ok=True) os.makedirs(test_folder, exist_ok=True) # Lấy danh sách các file trong thư mục hiện tại files = os.listdir(folder_path) random.shuffle(files) # Tính số lượng file cho train và test split_index = int(len(files) * split_ratio) train_files = files[:split_index] test_files = files[split_index:] # Di chuyển file vào thư mục train for file in train_files: src_file = os.path.join(folder_path, file) dst_file = os.path.join(train_folder, file) shutil.copy2(src_file, dst_file) # Di chuyển file vào thư mục test for file in test_files: src_file = os.path.join(folder_path, file) dst_file = os.path.join(test_folder, file) shutil.copy2(src_file, dst_file) source_dir = "A" # Thư mục chính chứa các thư mục con train_dir = "A_train" # Thư mục để lưu tập huấn luyện test_dir = "A_test" # Thư mục để lưu tập kiểm thử split_ratio = 0.8 # Tỷ lệ chia, ví dụ 0.8 tức là 80% train và 20% test split_data(source_dir, train_dir, test_dir, split_ratio)
Leave a Comment