File size: 1,354 Bytes
361b108 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | import pandas as pd
from sklearn.model_selection import StratifiedKFold
csv_path = '/workspace/odelia-data-scripts/csvs/combined_metadata.csv'
df = pd.read_csv(csv_path)
# Ensure Lesion is integer
df['Lesion'] = df['Lesion'].astype(int)
# Step 1: Get max lesion label per patient
patient_labels = df.groupby('PatientID')['Lesion'].max().reset_index()
# Step 2: Prepare data for stratified k-fold
patient_ids = patient_labels['PatientID'].values
labels = patient_labels['Lesion'].values
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
# Loop through folds
for fold, (train_idx, val_idx) in enumerate(skf.split(patient_ids, labels), 1):
train_patients = patient_ids[train_idx]
val_patients = patient_ids[val_idx]
# Filter rows for the current split
train_df = df[df['PatientID'].isin(train_patients)][['UID', 'Lesion', 'source']]
val_df = df[df['PatientID'].isin(val_patients)][['UID', 'Lesion', 'source']]
# Save to CSV
train_df.rename(columns={'UID': 'patient_id', 'Lesion': 'label'}).to_csv(
f'/workspace/odelia-data-scripts/odelia/train_ids_fold{fold}.csv', index=False
)
val_df.rename(columns={'UID': 'patient_id', 'Lesion': 'label'}).to_csv(
f'/workspace/odelia-data-scripts/odelia/val_ids_fold{fold}.csv', index=False
)
print("5-fold cross-validation splits saved.")
|