
#################################################
### THIS FILE WAS AUTOGENERATED! DO NOT EDIT! ###
#################################################
# file to edit: dev_nb/datablock.ipynb

from exp.nb_packages import *

train = pd.read_csv('train_with_strat_folds.csv')
DATA_PATH = Path('data')
IMG_PATH = DATA_PATH / 'images'
LABEL_COLS = ['healthy', 'multiple_diseases', 'rust', 'scab']

def get_data(SEED, BS, IMG_SIZE, fold):
    train_df_no_val = train.query(f'fold != {fold}')
    train_df_just_val = train.query(f'fold == {fold}')

    train_df_bal = pd.concat(
        [train_df_no_val.query('label != "multiple_diseases"'), train_df_just_val] +
        [train_df_no_val.query('label == "multiple_diseases"')] * 2
    ).sample(frac=1.0, random_state=SEED).reset_index(drop=True)

    datablock = DataBlock(
        blocks=(ImageBlock, CategoryBlock(vocab=LABEL_COLS)),
        getters=[
            ColReader('image_id', pref=IMG_PATH, suff='.jpg'),
            ColReader('label')
        ],
        splitter=IndexSplitter(train_df_bal.loc[train_df_bal.fold==fold].index),
        item_tfms=Resize(IMG_SIZE),
        batch_tfms=aug_transforms(size=IMG_SIZE, max_rotate=30., min_scale=0.75, flip_vert=True, do_flip=True)
    )
    return datablock.dataloaders(source=train_df_bal, bs=BS)