{"cells":[{"metadata":{},"cell_type":"markdown","source":"In this kernel, I want to show how to do stratified 5fold split of data. This is inspired by this [kernel](https://www.kaggle.com/lightforever/severstal-mlcomp-catalyst-train-0-90672-offline) from the recently concluded Steel Competition. Then we will build a dataloader in Pytorch showing how to do 5-fold training.  "},{"metadata":{},"cell_type":"markdown","source":"Fist we will need to install `mlcomp` since this library has nice inbuilt `stratified_group_k_fold` module for stratified K-folding. "},{"metadata":{"trusted":true,"_kg_hide-output":true},"cell_type":"code","source":"!pip install mlcomp","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torchvision\nimport torch.nn.functional as F\nimport torch.utils.data as data\nfrom collections import Counter\nimport albumentations\nfrom albumentations.pytorch import ToTensor\nfrom mlcomp.contrib.split import stratified_group_k_fold","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls ../input/understanding_cloud_organization","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train = pd.read_csv('../input/understanding_cloud_organization/train.csv')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> First let's create a column called `exists` which will indicate whether the image has mask(s) or not."},{"metadata":{"trusted":true},"cell_type":"code","source":"train['exists'] = train['EncodedPixels'].notnull().astype(int)\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> Now let's split the `Image_Label` to get the `image_name` and the corresponding `class_name`. Since the `Label` is a class name(`string`), we will map the `string` to `int` indicating the corresponding class labels; this will be our `class_id` "},{"metadata":{"trusted":true},"cell_type":"code","source":"train['image_name'] = train['Image_Label'].map(lambda x: x.split('_')[0].strip())\ntrain['class_name'] = train['Image_Label'].map(lambda x: x.split('_')[-1])\nclass_names_dict = {'Fish':1, 'Flower':2, 'Gravel':3, 'Sugar':4}\ntrain['class_id'] = train['class_name'].map(class_names_dict)\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> We will now zeroing out the corresponding values in `class_id` for which there exists no mask(s) i.e. `exists == 0`"},{"metadata":{"trusted":true},"cell_type":"code","source":"train['class_id'] = [row.class_id if row.exists else 0 for row in train.itertuples()]\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Finally we will do stratified 5-fold split based on `class_id` such that in each fold there are enough images from each class. This is where we will use the `stratified_group_k_fold` from `mlcomp`"},{"metadata":{"trusted":true},"cell_type":"code","source":"# You can change n_splits to any number you like. 5-fold split is the most common\ntrain['fold'] = stratified_group_k_fold(label='class_id', group_column='image_name', df=train, n_splits=5)\ntrain.head()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Let's check how many images are there in each from different classes"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"for fold in range(5):\n    print('-'*10, f'fold: {fold}', '-'*10)\n    df_fold = train[train['fold']==fold]\n    print('Images per class: ', Counter(df_fold['class_id']))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> It looks pretty OK to me. What do you think? Let's save it then!!"},{"metadata":{"trusted":true},"cell_type":"code","source":"train.to_csv('df_5fold.csv', index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> Now, we will use this `train` dataframe to build our fold data-loader in `pytorch`. First we will need some-helper function from this  [kernel](https://www.kaggle.com/ryches/turbo-charging-andrew-s-pytorch/notebook#Setting-up-data-for-training-in-Catalyst)."},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_mask(df: pd.DataFrame, image_name: str='img.jpg', shape: tuple = (350, 525)):\n    \"\"\"\n    Create mask based on df, image name and shape.\n    \"\"\"\n    masks = np.zeros((shape[0], shape[1], 4), dtype=np.float32)\n    df = df[df[\"image_name\"] == image_name]\n    for idx, im_name in enumerate(df[\"image_name\"].values):\n        for classidx, classid in enumerate([\"Fish\", \"Flower\", \"Gravel\", \"Sugar\"]):\n            mask = cv2.imread(\"../input/understanding-clouds-resized/train_masks_525/train_masks_525/\" + classid + im_name)\n            if mask is None:\n                continue\n            if mask[:,:,0].shape != (350,525):\n                mask = cv2.resize(mask, (525,350))\n            masks[:, :, classidx] = mask[:,:,0]\n    masks = masks/255\n    return masks","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ls ../input/understanding_cloud_organization","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Cloud_Dataset(data.Dataset):\n    def __init__(self, df, mode, transform=None, fold_index=None):\n        \n        self.df = df\n        self.transform = transform\n        self.mode = mode\n        \n        # change to your path\n        self.train_image_path = r'../input/understanding-clouds-resized/train_images_525/train_images_525/'\n        self.test_image_path = r'../input/understanding_cloud_organization/test_images'\n\n        self.fold_index = None\n        self.set_mode(mode, fold_index)\n\n    def set_mode(self, mode, fold_index):\n        self.mode = mode\n        self.fold_index = fold_index\n\n        if self.mode == 'train':\n            self.df_fold = self.df[self.df.fold != fold_index]\n            \n            self.img_ids = self.df_fold.image_name.values.tolist()\n            self.defects = self.df_fold.class_id.values.tolist()\n            self.exist_labels = self.df_fold.exists.astype(bool).values\n\n            self.num_data = len(self.df_fold)\n\n        elif self.mode == 'valid':\n            self.df_fold = self.df[self.df.fold == fold_index]\n            \n            self.img_ids = self.df_fold.image_name.values.tolist()\n            self.defects = self.df_fold.class_id.values.tolist()\n            self.exist_labels = self.df_fold.exists.astype(bool).values\n\n            self.num_data = len(self.df_fold)\n\n        elif self.mode == 'test':\n            self.test_list = sorted(os.listdir(self.test_image_path))\n            self.num_data = len(self.test_list)\n\n    def __getitem__(self, index):\n        if self.fold_index is None and self.mode != 'test':\n            print('WRONG!!!!!!! fold index is NONE!!!!!!!!!!!!!!!!!')\n            return\n        \n        if self.mode == 'test':\n            image = cv2.imread(os.path.join(self.test_image_path, self.test_list[index]), 1)\n            if self.transform:\n                sample = {\"image\": image}\n                sample = self.transform(**sample)\n                image = sample['image']\n            image_id = self.test_list[index].replace('.png', '')\n            return image_id, image\n        \n        elif self.mode != 'test':\n            image_id = self.img_ids[index]\n            mask = make_mask(self.df_fold, image_id)\n            image = cv2.imread(os.path.join(self.train_image_path, image_id), 1)\n            \n        if self.transform:\n            augmented = self.transform(image=image, mask=mask)\n            image = augmented['image']\n            mask = augmented['mask'] # 1x320x320x4\n            mask = mask[0].permute(2, 0, 1) # 1x4x320x320\n            \n        return image, mask\n             \n    def __len__(self):\n        return self.num_data","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> We will also need a function for generating our `transformation` for `train`, `valid` and `test` loaders."},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef generate_transforms(mode):\n    # MAX_SIZE = 448\n    IMAGE_SIZE = [320,320]\n\n    train_transform = albumentations.Compose([\n        \n        albumentations.HorizontalFlip(p=1),\n        albumentations.ElasticTransform(p=1, alpha=120, sigma=120 * 0.05, alpha_affine=120 * 0.03),\n        albumentations.Resize(IMAGE_SIZE[0], IMAGE_SIZE[1]),\n        albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0),\n        ToTensor()\n    \n    ])\n\n\n    val_transform = albumentations.Compose([\n        albumentations.Resize(IMAGE_SIZE[0], IMAGE_SIZE[1]),\n        albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0),\n        ToTensor()\n    \n    ])\n\n    test_transform = albumentations.Compose([\n        #albumentations.Resize(IMAGE_SIZE[0], IMAGE_SIZE[1]),\n        albumentations.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225), max_pixel_value=255.0, p=1.0),\n        ToTensor()\n    \n    ])\n    \n    \n    if mode == 'train':\n        return train_transform\n    elif mode == 'valid':\n        return val_transform\n    else:\n        return test_transform","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"\n\ndef get_fold_dataloader(fold_index, mode, batch_size = 16):\n    df = pd.read_csv('df_5fold.csv')\n    \n    dataset = Cloud_Dataset(df, mode, generate_transforms(mode), fold_index) # df, mode, transform=None, fold_index=None\n    \n    dataloader = torch.utils.data.DataLoader(dataset,\n                            batch_size=batch_size,\n                            num_workers=0,\n                            pin_memory=True,\n                            shuffle=mode == 'train',\n                            )\n    \n    return dataloader\n\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> Let's get our `train_loader`, `valid_loader` and `test_loader`"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_loader = get_fold_dataloader(fold_index=0, mode='train')\nvalid_loader = get_fold_dataloader(fold_index=0, mode='valid')\ntest_loader = get_fold_dataloader(fold_index=0, mode='test')","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> Time to check if they are working propely as expected or not!"},{"metadata":{"trusted":true},"cell_type":"code","source":"for image, mask in train_loader:\n    print(image.shape, mask.shape)\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for image, mask in valid_loader:\n    print(image.shape, mask.shape)\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for image_id, image in test_loader:\n    print(image.shape)\n    break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":">I guess our loader are working fine as expected. This is my way of creating loaders. Let me know your way?"},{"metadata":{},"cell_type":"markdown","source":"Below a pseudo-code for how to do 5-fold training!\n```python\nfor fold in range(num_folds=5):\n    train_loader = get_fold_dataloader(fold_index=fold, mode='train')\n    valid_loader = get_fold_dataloader(fold_index=fold, mode='valid')\n    test_loader = get_fold_dataloader(fold_index=fold, mode='test')\n    \n    train_model(model, train_loader)\n    validate_model(model, valid_loader)\n    test_model(model, test_loader)\n    save_model(model, fold)\n```\n\n"},{"metadata":{},"cell_type":"markdown","source":"> I hope this kernel was useful to you in someway. If you have any suggestions/queries, please let me know in the comments."}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":1}