application_vnd.jupyter.stderr_output_1.png
import pandas as pd

NUM_FOLDS = 5
import pandas as pd
from tqdm import tqdm
data_dir = '../input/ranzcr-clip-catheter-line-classification/'
target_col = ['ETT - Abnormal', 'ETT - Borderline', 'ETT - Normal', 'NGT - Abnormal', 'NGT - Borderline', 'NGT - Incompletely Imaged', 'NGT - Normal', 'CVC - Abnormal', 'CVC - Borderline', 'CVC - Normal', 'Swan Ganz Catheter Present']
def annotation_col_name(x):
    x = x[1]
    return 'annotation-' + ''.join([c for c in x if c.upper() == c and c != ' '])
annotations_df = pd.read_csv(data_dir + 'train_annotations.csv', usecols=['StudyInstanceUID', 'label'])
annotations_df['ones'] = 1
annotations_df = pd.pivot_table(annotations_df, values=['ones'], index=['StudyInstanceUID'], columns=['label'], aggfunc=np.sum).fillna(0)
annotations_df.columns = [annotation_col_name(col_name) for col_name in annotations_df.columns]
train_gt_df = pd.read_csv(data_dir + '/train.csv')
train_gt_df = train_gt_df.merge(annotations_df, on='StudyInstanceUID', how='left').fillna(0)
target_col += [x for x in annotations_df.columns]
num_target = len(target_col)
kfold_df = train_gt_df.copy().drop(['StudyInstanceUID'], axis=1)

patient_kfold_df = kfold_df.groupby('PatientID').sum()
priority_df = patient_kfold_df.sum().to_frame().reset_index().rename(columns={'index': 'label', 0: 'score'})
priority_df['score'] = [s * 4 if 'annotation' in l else s for l, s in zip(priority_df['label'], priority_df['score'])]
priority_df.sort_values(by='score', inplace=True)
label_priority = list(priority_df['label'])
label_importance = [(0.1 + patient_kfold_df.sum().max() - x) ** 2 for x in priority_df['score']]
label_importance[0] = 3 * label_importance[0]
patient_kfold_df = patient_kfold_df.sort_values(by=label_priority, ascending=False)
def kfold_quality(folds):
    patient_kfold_df['fold'] = folds + [-1] * (len(patient_kfold_df) - len(folds))
    scores = patient_kfold_df[:len(folds)][['fold'] + target_col].groupby('fold').sum()
    if scores.shape[0] < NUM_FOLDS:
        scores.reset_index(inplace=True)
        for f in range(NUM_FOLDS):
            if f not in scores.index:
                scores.loc[len(scores)] = 0
                scores.loc[len(scores) - 1, 'fold'] = f
        scores.set_index('fold', drop=True, inplace=True)
    score = scores.sum(axis=1).std()
    for label, importance in zip(label_priority, label_importance):
        scores[label] *= importance
        score += scores[label].max() - scores[label].min()
    return score

patient_kfold_df = patient_kfold_df.reset_index()
train_gt_df = train_gt_df.merge(patient_kfold_df[['PatientID', 'fold']], on='PatientID', how='left')
patient_kfold_df[['PatientID', 'fold']].to_csv(f'stratified_{NUM_FOLDS}_folds_by_patient_id.csv', index=False)
train_gt_df[['StudyInstanceUID', 'fold']].to_csv(f'stratified_{NUM_FOLDS}_folds.csv', index=False)

patients_per_fold = train_gt_df[['fold', 'PatientID']].groupby('fold').agg({'PatientID': 'nunique'}).rename(columns={'PatientID': 'Num. patients'})
patients_per_fold