age_binning.py

import pandas as pd
import numpy as np
import os, nibabel, shutil

from collections import Counter
from sklearn.model_selection import train_test_split

data = pd.read_csv("masterdata.csv")
age_dict, age_dict_rounded = {}, {}

for index, row in data.iterrows():
    age_dict[row['patientid']] = float(row['age'])

for pid in age_dict:
    age_dict_rounded[pid] = round(age_dict[pid])

# age = list(age_dict.values())
# rounded_age = list(age_dict_rounded.values())
# print(age_dict_rounded)

X, y = [], []
for patient_id in age_dict_rounded:
    X.append(patient_id)
    y.append(age_dict_rounded[patient_id])

X = np.array(X)
y = np.array(y)

X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y, test_size=0.10)
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, stratify=y_train, test_size=0.10)

print(len(X), len(X_train), len(X_val), len(X_test))

if os.path.exists(f"./data"):
    shutil.rmtree(f"./data")
    print(f"removed an existing ./data directory")

for x in ["test", "val", "train"]:
    if not os.path.exists(f"./data/HC/{x}"):
        os.makedirs(f"./data/HC/{x}")
    for z in set(y_train):
        if not os.path.exists(f"./data/HC/{x}/{z}"):
            os.makedirs(f"./data/HC/{x}/{z}")

for x in X_test:
    name = x.split("/")[-1]
    img = nibabel.load(x)
    try: 
        data = img.get_fdata()
        os.symlink(x, f"./data/HC/test/{age_dict_rounded[x]}/{name}")
    except:
        print("Corrupted file")
print("Done with test")
for x in X_val:
    name = x.split("/")[-1]
    img = nibabel.load(x)
    try: 
        data = img.get_fdata()
        os.symlink(x, f"./data/HC/val/{age_dict_rounded[x]}/{name}")
    except:
        print("Corrupted file")
print("Done with val")
for x in X_train:
    name = x.split("/")[-1]
    img = nibabel.load(x)
    try: 
        data = img.get_fdata()
        os.symlink(x, f"./data/HC/train/{age_dict_rounded[x]}/{name}")
    except:
        print("Corrupted file")
print("Done with train")