Untitled

mail@pastecode.io avatar
unknown
plain_text
14 days ago
2.4 kB
3
Indexable
Never
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