Cataract-ViT / src /prepare_dataset.py
Decoder24's picture
Upload folder using huggingface_hub
a080b32 verified
# prepare_dataset.py
import os
import shutil
import pandas as pd
from sklearn.model_selection import train_test_split
# --- KONFIGURASI ---
# Path ke file CSV yang sudah lengkap labelnya
CSV_PATH = 'outputs/csv/labels.csv'
# Path ke dataset asli yang berisi folder 01, 02, ...
SOURCE_DATA_DIR = 'data/Nuclear Cataract Database for Biomedical and Machine Learning Applications/Nuclear Cataract Dataset'
# Path ke folder baru tempat data split akan disimpan
SPLIT_DATA_DIR = 'data/split data'
# Rasio untuk data validasi (0.2 berarti 20%)
VALIDATION_SPLIT_RATIO = 0.2
# --- SCRIPT UTAMA ---
print("Membaca file metadata labels.csv...")
df = pd.read_csv(CSV_PATH)
# Ekstrak ID pasien dari path gambar
# Misal: dari '01/DER/image_01.png' menjadi '01'
df['patient_id'] = df['image'].apply(lambda x: x.split('/')[0])
# Dapatkan daftar unik semua pasien
unique_patients = df['patient_id'].unique()
print(f"Total pasien ditemukan: {len(unique_patients)}")
# Bagi daftar pasien menjadi train dan validation
train_patients, valid_patients = train_test_split(
unique_patients,
test_size=VALIDATION_SPLIT_RATIO,
random_state=42 # random_state untuk hasil yang bisa direproduksi
)
print(f"Jumlah pasien untuk training: {len(train_patients)}")
print(f"Jumlah pasien untuk validasi: {len(valid_patients)}")
# Fungsi untuk menyalin file
def copy_files(patient_list, target_folder):
target_path = os.path.join(SPLIT_DATA_DIR, target_folder)
subset_df = df[df['patient_id'].isin(patient_list)]
print(f"\nMemproses {len(subset_df)} gambar untuk set '{target_folder}'...")
for index, row in subset_df.iterrows():
grade = str(row['grade'])
image_path = row['image']
class_folder = os.path.join(target_path, grade)
os.makedirs(class_folder, exist_ok=True)
source_file = os.path.join(SOURCE_DATA_DIR, image_path)
destination_file = os.path.join(class_folder, os.path.basename(image_path))
# Cek apakah file sumber ada sebelum menyalin
if os.path.exists(source_file):
shutil.copyfile(source_file, destination_file)
else:
print(f" Peringatan: File sumber tidak ditemukan -> {source_file}")
# Buat folder utama untuk data split
if os.path.exists(SPLIT_DATA_DIR):
shutil.rmtree(SPLIT_DATA_DIR)
os.makedirs(SPLIT_DATA_DIR)
# Jalankan proses penyalinan
copy_files(train_patients, 'train')
copy_files(valid_patients, 'valid')
print("\n--- Proses splitting per pasien selesai! ---")
print(f"Dataset baru siap di folder '{SPLIT_DATA_DIR}' dengan struktur yang benar.")