create_splits.py

import pdb
import os
import pandas as pd
from datasets.dataset_mtl_concat import Generic_WSI_MTL_Dataset, Generic_MIL_MTL_Dataset, save_splits
import argparse
import numpy as np

parser = argparse.ArgumentParser(description='Creating splits for whole slide classification')
parser.add_argument('--label_frac', type=float, default= -1,
										help='fraction of labels (default: [1.0])')
parser.add_argument('--seed', type=int, default=1,
										help='random seed (default: 1)')
parser.add_argument('--k', type=int, default=10,
										help='number of splits (default: 10)')
parser.add_argument('--hold_out_test', action='store_true', default=False,
										help='fraction to hold out (default: 0)')
parser.add_argument('--split_code', type=str, default=None)
parser.add_argument('--task', type=str, choices=['dummy_mtl_concat'])

args = parser.parse_args()

if args.task == 'dummy_mtl_concat':
    args.n_classes=18
    dataset = Generic_WSI_MTL_Dataset(csv_path = 'dataset_csv/dummy_dataset.csv',
                            shuffle = False, 
                            seed = args.seed, 
                            print_info = True,
                            label_dicts = [{'Lung':0, 'Breast':1, 'Colorectal':2, 'Ovarian':3, 
                                                                'Pancreatic':4, 'Adrenal':5, 
                                                                'Skin':6, 'Prostate':7, 'Renal':8, 'Bladder':9, 
                                                                'Esophagagostric':10,  'Thyroid':11,
                                                                'Head Neck':12,  'Glioma':13, 
                                                                'Germ Cell':14, 'Endometrial': 15, 'Cervix': 16, 'Liver': 17},
                                            {'Primary':0,  'Metastatic':1},
                                            {'F':0, 'M':1}],
                            label_cols = ['label', 'site', 'sex'],
                            patient_strat= False)

         
else:
	raise NotImplementedError

num_slides_cls = np.array([len(cls_ids) for cls_ids in dataset.patient_cls_ids])
val_num = np.floor(num_slides_cls * 0.1).astype(int)
test_num = np.floor(num_slides_cls * 0.2).astype(int)

print(val_num)
print(test_num)

if __name__ == '__main__':
		if args.label_frac > 0:
			label_fracs = [args.label_frac]
		else:
			label_fracs = [1.0]

		if args.hold_out_test:
			custom_test_ids = dataset.sample_held_out(test_num=test_num)
		else:
			custom_test_ids = None
		for lf in label_fracs:
			if args.split_code is not None:
				split_dir = 'splits/'+ str(args.split_code) + '_{}'.format(int(lf * 100))
			else:
				split_dir = 'splits/'+ str(args.task) + '_{}'.format(int(lf * 100))
			
			dataset.create_splits(k = args.k, val_num = val_num, test_num = test_num, label_frac=lf, custom_test_ids=custom_test_ids)

			os.makedirs(split_dir, exist_ok=True)
			for i in range(args.k):
				if dataset.split_gen is None:
					ids = []
					for split in ['train', 'val', 'test']:
						ids.append(dataset.get_split_from_df(pd.read_csv(os.path.join(split_dir, 'splits_{}.csv'.format(i))), split_key=split, return_ids_only=True))
					
					dataset.train_ids = ids[0]
					dataset.val_ids = ids[1]
					dataset.test_ids = ids[2]
				else:
					dataset.set_splits()

				descriptor_df = dataset.test_split_gen(return_descriptor=True)
				descriptor_df.to_csv(os.path.join(split_dir, 'splits_{}_descriptor.csv'.format(i)))
				
				splits = dataset.return_splits(from_id=True)
				save_splits(splits, ['train', 'val', 'test'], os.path.join(split_dir, 'splits_{}.csv'.format(i)))
				save_splits(splits, ['train', 'val', 'test'], os.path.join(split_dir, 'splits_{}_bool.csv'.format(i)), boolean_style=True)