{"cells":[{"metadata":{},"cell_type":"markdown","source":"# An Implementation of GridMask Based on albumentations\n\n\nHi, here is my implementation of GridMask augmentation based on albumentations. If you find it helpful please upvote me.\nThanks!\n\nGridMask: https://arxiv.org/abs/2001.04086\n\n![](https://www.googleapis.com/download/storage/v1/b/kaggle-user-content/o/inbox%2F448347%2F2184811fc18555da73a64277ab5016a1%2F2020-01-30%2011.33.43.png?generation=1580351647542803&alt=media)\n\nalbumentations： https://github.com/albumentations-team/albumentations"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport albumentations\nfrom albumentations.core.transforms_interface import DualTransform\nfrom albumentations.augmentations import functional as F\nimport matplotlib.pyplot as plt\nimport torch\nfrom torch.utils.data import TensorDataset, DataLoader, Dataset\n\ndata_dir = '../input/bengaliai-cv19'\nfiles_train = [f'train_image_data_{fid}.parquet' for fid in range(1)]\n\nHEIGHT = 137\nWIDTH = 236","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"df_train = pd.read_csv(os.path.join(data_dir, f'train.csv'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def read_data(files):\n    tmp = []\n    for f in files:\n        F = os.path.join(data_dir, f)\n        data = pd.read_parquet(F)\n        tmp.append(data)\n    tmp = pd.concat(tmp)\n\n    data = tmp.iloc[:, 1:].values\n    return data\n\n# train data\ndata_train = read_data(files_train)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GridMask(DualTransform):\n    \"\"\"GridMask augmentation for image classification and object detection.\n    \n    Author: Qishen Ha\n    Email: haqishen@gmail.com\n    2020/01/29\n\n    Args:\n        num_grid (int): number of grid in a row or column.\n        fill_value (int, float, lisf of int, list of float): value for dropped pixels.\n        rotate ((int, int) or int): range from which a random angle is picked. If rotate is a single int\n            an angle is picked from (-rotate, rotate). Default: (-90, 90)\n        mode (int):\n            0 - cropout a quarter of the square of each grid (left top)\n            1 - reserve a quarter of the square of each grid (left top)\n            2 - cropout 2 quarter of the square of each grid (left top & right bottom)\n\n    Targets:\n        image, mask\n\n    Image types:\n        uint8, float32\n\n    Reference:\n    |  https://arxiv.org/abs/2001.04086\n    |  https://github.com/akuxcw/GridMask\n    \"\"\"\n\n    def __init__(self, num_grid=3, fill_value=0, rotate=0, mode=0, always_apply=False, p=0.5):\n        super(GridMask, self).__init__(always_apply, p)\n        if isinstance(num_grid, int):\n            num_grid = (num_grid, num_grid)\n        if isinstance(rotate, int):\n            rotate = (-rotate, rotate)\n        self.num_grid = num_grid\n        self.fill_value = fill_value\n        self.rotate = rotate\n        self.mode = mode\n        self.masks = None\n        self.rand_h_max = []\n        self.rand_w_max = []\n\n    def init_masks(self, height, width):\n        if self.masks is None:\n            self.masks = []\n            n_masks = self.num_grid[1] - self.num_grid[0] + 1\n            for n, n_g in enumerate(range(self.num_grid[0], self.num_grid[1] + 1, 1)):\n                grid_h = height / n_g\n                grid_w = width / n_g\n                this_mask = np.ones((int((n_g + 1) * grid_h), int((n_g + 1) * grid_w))).astype(np.uint8)\n                for i in range(n_g + 1):\n                    for j in range(n_g + 1):\n                        this_mask[\n                             int(i * grid_h) : int(i * grid_h + grid_h / 2),\n                             int(j * grid_w) : int(j * grid_w + grid_w / 2)\n                        ] = self.fill_value\n                        if self.mode == 2:\n                            this_mask[\n                                 int(i * grid_h + grid_h / 2) : int(i * grid_h + grid_h),\n                                 int(j * grid_w + grid_w / 2) : int(j * grid_w + grid_w)\n                            ] = self.fill_value\n                \n                if self.mode == 1:\n                    this_mask = 1 - this_mask\n\n                self.masks.append(this_mask)\n                self.rand_h_max.append(grid_h)\n                self.rand_w_max.append(grid_w)\n\n    def apply(self, image, mask, rand_h, rand_w, angle, **params):\n        h, w = image.shape[:2]\n        mask = F.rotate(mask, angle) if self.rotate[1] > 0 else mask\n        mask = mask[:,:,np.newaxis] if image.ndim == 3 else mask\n        image *= mask[rand_h:rand_h+h, rand_w:rand_w+w].astype(image.dtype)\n        return image\n\n    def get_params_dependent_on_targets(self, params):\n        img = params['image']\n        height, width = img.shape[:2]\n        self.init_masks(height, width)\n\n        mid = np.random.randint(len(self.masks))\n        mask = self.masks[mid]\n        rand_h = np.random.randint(self.rand_h_max[mid])\n        rand_w = np.random.randint(self.rand_w_max[mid])\n        angle = np.random.randint(self.rotate[0], self.rotate[1]) if self.rotate[1] > 0 else 0\n\n        return {'mask': mask, 'rand_h': rand_h, 'rand_w': rand_w, 'angle': angle}\n\n    @property\n    def targets_as_params(self):\n        return ['image']\n\n    def get_transform_init_args_names(self):\n        return ('num_grid', 'fill_value', 'rotate', 'mode')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class BengaliDataset(Dataset):\n    def __init__(self, csv, data, idx, split, mode, image_size, transform=None):\n\n        self.csv = csv.reset_index()\n        self.data = data\n        self.idx = np.asarray(idx)\n        self.split = split\n        self.mode = mode\n        self.image_size = image_size\n        self.transform = transform\n\n    def __len__(self):\n        return self.idx.shape[0]\n\n    def __getitem__(self, index):\n        index = self.idx[index]\n        this_img_id = self.csv.iloc[index].image_id\n        \n        image = self.data[index].reshape(HEIGHT, WIDTH)\n        image = cv2.resize(image, (self.image_size, self.image_size))\n\n        if self.transform is not None:\n            res = self.transform(image=image)\n            image = res['image'].astype(np.float32)\n        else:\n            image = image.astype(np.float32)\n\n        image /= 255\n        image = image[np.newaxis, :, :]\n#         image = np.repeat(image, 3, 0)  # 1ch to 3ch\n\n        if self.mode == 'test':\n            return torch.tensor(image)\n        else:\n            label_1 = self.csv.iloc[index].grapheme_root\n            label_2 = self.csv.iloc[index].vowel_diacritic\n            label_3 = self.csv.iloc[index].consonant_diacritic\n            label = [label_1, label_2, label_3]\n            return torch.tensor(image), torch.tensor(label)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_imgs(dataset_show):\n    from pylab import rcParams\n    rcParams['figure.figsize'] = 20,10\n    for i in range(2):\n        f, axarr = plt.subplots(1,5)\n        for p in range(5):\n            idx = np.random.randint(0, len(dataset_show))\n            img, label = dataset_show[idx]\n            axarr[p].imshow(img.transpose(0, 1).transpose(1,2).squeeze())\n            axarr[p].set_title(idx)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"markdown","source":"# Usage Example\n\n### num_grid = 3"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    GridMask(num_grid=3, p=1),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"\n### num_grid = (3,7)"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    GridMask(num_grid=(3,7), p=1),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### num_grid = 3, rotate = 15"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    GridMask(num_grid=3, rotate=15, p=1),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### num_grid = (3,7), mode = 1"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    GridMask(num_grid=(3,7), mode=1, p=1),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### num_grid = 3, mode = 2"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    GridMask(num_grid=3, mode=2, p=1),\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Combine mode 0,1,2"},{"metadata":{"trusted":true},"cell_type":"code","source":"transforms_train = albumentations.Compose([\n    albumentations.OneOf([\n        GridMask(num_grid=3, mode=0),\n        GridMask(num_grid=3, mode=1),\n        GridMask(num_grid=3, mode=2),\n    ], p=1)\n])\n\ndf_show = df_train.iloc[:1000]\ndataset_show = BengaliDataset(df_show, data_train, list(range(df_show.shape[0])), 'train', 'train', 128, transform=transforms_train)\nplot_imgs(dataset_show)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"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}