Download and split ixi

import hashlib
import os
import requests
from tqdm import tqdm
import argparse
from zipfile import ZipFile
import shutil
import pandas as pd
from sklearn.model_selection import train_test_split

from fedbiomed.node.config import NodeConfig 
from fedbiomed.node.environ import environ

def parse_args():
    parser = argparse.ArgumentParser(description='IXI Sample downloader and splitter')
    parser.add_argument('-f', '--root_folder', required=True, type=str)
    parser.add_argument('-F', '--force',  action=argparse.BooleanOptionalAction, required=False, type=bool, default=False)
    return parser.parse_args()

def has_correct_checksum_md5(filename, hash):
    with open(filename, "rb") as f:
        file_hash = hashlib.md5()
        while chunk := f.read(8192):
            file_hash.update(chunk)
    return str(file_hash.hexdigest()) == hash

def download_file(url, filename):
    """
    Helper method handling downloading large files from `url` to `filename`. Returns a pointer to `filename`.
    """
    print('Downloading file from:', url)
    print('File will be saved as:', filename)
    chunkSize = 1024
    r = requests.get(url, stream=True)
    with open(filename, 'wb') as f:
        pbar = tqdm(unit="B", total=int(r.headers['Content-Length']))
        for chunk in r.iter_content(chunk_size=chunkSize):
            if chunk:  # filter out keep-alive new chunks
                pbar.update(len(chunk))
                f.write(chunk)
    return filename

def download_and_extract_ixi_sample(root_folder):
    url = 'https://prod-dcd-datasets-cache-zipfiles.s3.eu-west-1.amazonaws.com/7kd5wj7v7p-3.zip'
    zip_filename = os.path.join(root_folder, 'notebooks', 'data', '7kd5wj7v7p-3.zip')
    data_folder = os.path.join(root_folder, 'notebooks', 'data')
    extracted_folder = os.path.join(data_folder, '7kd5wj7v7p-3', 'IXI_sample')

    # Extract if ZIP exists but not folder
    if not os.path.exists(zip_filename):
        # Download if it does not exist
        download_file(url, zip_filename)
        
    # Check if extracted folder exists
    if os.path.isdir(extracted_folder):
        print(f'Dataset folder already exists in {extracted_folder}')
        return extracted_folder

    assert has_correct_checksum_md5(zip_filename, 'eecb83422a2685937a955251fa45cb03')
    with ZipFile(zip_filename, 'r') as zip_obj:
        zip_obj.extractall(data_folder)

    assert os.path.isdir(extracted_folder)
    return extracted_folder

if __name__ == '__main__':
    args = parse_args()
    root_folder = os.path.abspath(os.path.expanduser(args.root_folder))
    assert os.path.isdir(root_folder), f'Folder does not exist: {root_folder}'

    # Centralized dataset
    centralized_data_folder = download_and_extract_ixi_sample(root_folder)

    # Federated Dataset
    federated_data_folder = os.path.join(root_folder, 'notebooks', 'data', 'Hospital-Centers')
    shutil.rmtree(federated_data_folder, ignore_errors=True)

    csv_global = os.path.join(centralized_data_folder, 'participants.csv')
    allcenters = pd.read_csv(csv_global)

    # Split centers
    center_names = ['Guys', 'HH', 'IOP']
    center_dfs = list()

    for center_name in center_names:
        cfg_folder = os.path.join(args.root_folder, 'etc')
        os.makedirs(cfg_folder, exist_ok=True)
        cfg_file = os.path.join(cfg_folder, f'{center_name.lower()}.ini')

        print(f'Creating node at: {cfg_file}')
        config = NodeConfig(root=environ["ROOT_DIR"], name=f'{center_name.lower()}.ini', auto_generate=False)
        if config.is_config_existing() and not args.force:
            print(f"**Warning: config file {cfg_file} already exists. To overwrite, please specify `--force` option")
        else:
            config.generate(force=args.force)
            config.set('default', 'id', center_name)
            config.write()

        df = allcenters[allcenters.SITE_NAME == center_name]
        center_dfs.append(df)

        train, test = train_test_split(df, test_size=0.1, random_state=21)

        train_folder = os.path.join(federated_data_folder, center_name, 'train')
        holdout_folder = os.path.join(federated_data_folder, center_name, 'holdout')
        if not os.path.exists(train_folder):
            os.makedirs(train_folder)
        if not os.path.exists(holdout_folder):
            os.makedirs(holdout_folder)

        for subject_folder in train.FOLDER_NAME.values:
            shutil.copytree(
                src=os.path.join(centralized_data_folder, subject_folder),
                dst=os.path.join(train_folder, subject_folder),
                dirs_exist_ok=True
            )

        train_participants_csv = os.path.join(train_folder, 'participants.csv')
        train.to_csv(train_participants_csv)

        for subject_folder in test.FOLDER_NAME.values:
            shutil.copytree(
                src=os.path.join(centralized_data_folder, subject_folder),
                dst=os.path.join(holdout_folder, subject_folder),
                dirs_exist_ok=True
            )
        test.to_csv(os.path.join(holdout_folder, 'participants.csv'))

    print(f'Centralized dataset located at: {centralized_data_folder}')
    print(f'Federated dataset located at: {federated_data_folder}')

    print()
    print('Please add the data to your nodes executing and using the `ixi-train` tag:')
    for center_name in center_names:
        print(f'\t./scripts/fedbiomed_run node --config {center_name.lower()}.ini dataset add')

    print()
    print('Then start your nodes by executing:')
    for center_name in center_names:
        print(f'\t./scripts/fedbiomed_run node --config {center_name.lower()}.ini start')