{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install segmentation-models-pytorch","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport cv2\nimport sys\nimport collections\nimport segmentation_models_pytorch as smp\nimport albumentations as albu\n\nfrom torch.utils.data import DataLoader\nfrom glob import glob\nfrom os import path\nfrom PIL import Image\nfrom tqdm import tqdm\nfrom sklearn.model_selection import train_test_split\nfrom matplotlib import pyplot as plt\n\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Constants"},{"metadata":{"trusted":true},"cell_type":"code","source":"IMG_WIDTH = 256\nIMG_HEIGHT = 256\nNUM_CLASSES = 46\n\nBATCH_SIZE = 32\nN_WORKERS = 2\n\nroot_path_train = '/kaggle/input/imaterialist-fashion-2019-FGVC6/train'\ndf_path_train = '/kaggle/input/imaterialist-fashion-2019-FGVC6/train.csv'","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prepare DataFrames"},{"metadata":{"trusted":true},"cell_type":"code","source":"def create_one_represent_class(df_param):\n    v_c_df = df_param['CategoryId'].value_counts().reset_index()\n    one_represent = v_c_df.loc[v_c_df['CategoryId'] == 1, 'index'].tolist()\n    df_param.loc[df_param['CategoryId'].isin(one_represent), 'CategoryId'] = 'one_represent'\n    return df_param\n\ndef custom_train_test_split(df_param):\n    \n    df_param['CategoryId'] = df_param.ClassId.apply(lambda x: str(x).split(\"_\")[0])\n    \n    img_categ = train_df.groupby('ImageId')['CategoryId'].apply(list).reset_index()\n    img_categ['CategoryId'] = img_categ['CategoryId'].apply(lambda x: ' '.join(sorted(x)))\n    \n    img_categ = create_one_represent_class(img_categ)\n    \n    img_train, img_val  = train_test_split(img_categ, test_size=0.2, random_state=42, stratify=img_categ['CategoryId'])\n    \n    df_param = df_param.drop(columns='CategoryId')\n    \n    df_train = df_param[df_param['ImageId'].isin(img_train['ImageId'])].reset_index(drop=True)\n    df_val = df_param[df_param['ImageId'].isin(img_val['ImageId'])].reset_index(drop=True)\n    \n    return df_train, df_val","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(df_path_train)\ntrain_df.head()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df, val_df = custom_train_test_split(train_df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"def rle_decode(mask_rle, shape):\n    '''\n    mask_rle: run-length as string formated: [start0] [length0] [start1] [length1]... in 1d array\n    shape: (height,width) of array to return\n    Returns numpy array according to the shape, 1 - mask, 0 - background\n    '''\n    shape = (shape[1], shape[0])\n    s = mask_rle.split()\n    # gets starts & lengths 1d arrays\n    starts, lengths = [np.asarray(x, dtype=int) for x in (s[0::2], s[1::2])]\n    starts -= 1\n    # gets ends 1d array\n    ends = starts + lengths\n    # creates blank mask image 1d array\n    img = np.zeros(shape[0] * shape[1], dtype=np.uint8)\n    # sets mark pixles\n    for lo, hi in zip(starts, ends):\n        img[lo:hi] = 1\n    # reshape as a 2d mask image\n    return img.reshape(shape).T  # Needed to align to RLE direction\n\n    \nclass UnetDataset(torch.utils.data.Dataset):\n    def __init__(self, image_dir, df, height, width, augmentation=None, preprocessing=None):\n        \n        self.preprocessing = preprocessing\n        self.augmentation = augmentation\n        \n        self.image_dir = image_dir\n        self.df = df\n        \n        self.height = height\n        self.width = width\n        \n        self.image_info = collections.defaultdict(dict)\n        \n        self.df['CategoryId'] = self.df.ClassId.apply(lambda x: str(x).split(\"_\")[0])\n        self.num_classes = self.df['CategoryId'].nunique()\n        \n        temp_df = self.df.groupby('ImageId')['EncodedPixels', 'CategoryId'].agg(lambda x: list(x)).reset_index()\n        size_df = self.df.groupby('ImageId')['Height', 'Width'].mean().reset_index()\n        temp_df = temp_df.merge(size_df, on='ImageId', how='left')\n        \n        for index, row in tqdm(temp_df.iterrows(), total=len(temp_df)):\n            image_id = row['ImageId']\n            image_path = os.path.join(self.image_dir, image_id)\n            self.image_info[index][\"image_id\"] = image_id\n            self.image_info[index][\"image_path\"] = image_path\n            self.image_info[index][\"width\"] = self.width\n            self.image_info[index][\"height\"] = self.height\n            self.image_info[index][\"labels\"] = row[\"CategoryId\"]\n            self.image_info[index][\"orig_height\"] = row[\"Height\"]\n            self.image_info[index][\"orig_width\"] = row[\"Width\"]\n            self.image_info[index][\"annotations\"] = row[\"EncodedPixels\"]\n\n    def __getitem__(self, idx):\n        \n        img_path = self.image_info[idx][\"image_path\"]\n        \n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img,cv2.COLOR_BGR2RGB)\n        img = cv2.resize(img, (self.width, self.height))\n\n        info = self.image_info[idx]\n        \n        mask = np.zeros((self.width, self.height, self.num_classes))\n        labels = np.zeros(self.num_classes)\n        \n        for annotation, label in zip(info['annotations'], info['labels']):\n            cur_mask = rle_decode(annotation, (info['orig_height'], info['orig_width']))\n            mask[:, :, int(label)] += cv2.resize(cur_mask, (self.width, self.height))\n            labels[int(label)] = 1\n            \n        mask = (mask > 0.5).astype(np.float32)\n        \n        # apply augmentations\n        if self.augmentation is not None:\n            sample = self.augmentation(image=img, mask=mask)\n            img, mask = sample['image'], sample['mask']\n        \n        # apply preprocessing\n        if self.preprocessing is not None:\n            sample = self.preprocessing(image=img, mask=mask)\n            img, mask = sample['image'], sample['mask']\n            \n        return img, mask, labels\n\n    def __len__(self):\n        return len(self.image_info)\n\n    \ndef collate_function(batch):\n    image_array = torch.zeros((len(batch), batch[0][0].shape[0], batch[0][0].shape[1], batch[0][0].shape[2]))\n    mask_array = torch.zeros((len(batch), batch[0][1].shape[0], batch[0][1].shape[1], batch[0][1].shape[2]))\n    label_array = torch.zeros((len(batch), batch[0][2].shape[0]))\n    \n    for i in range(len(batch)):\n        image_array[i,:,:,:] = torch.Tensor(batch[i][0])\n        mask_array[i,:,:,:] = torch.Tensor(batch[i][1])\n        label_array[i,:] = torch.Tensor(batch[i][2])\n        \n    return image_array, (mask_array, label_array)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Utils"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_training_augmentation():\n    train_transform = [\n        albu.HorizontalFlip(p=0.5),\n        albu.VerticalFlip(p=0.5),\n    ]\n    return albu.Compose(train_transform)\n\n\ndef to_tensor(x, **kwargs):\n    return x.transpose(2, 0, 1).astype('float32')\n\ndef get_preprocessing(preprocessing_fn):\n    \"\"\"Construct preprocessing transform\n    \n    Args:\n        preprocessing_fn (callbale): data normalization function \n            (can be specific for each pretrained neural network)\n    Return:\n        transform: albumentations.Compose\n    \n    \"\"\"\n    \n    _transform = [\n        albu.Lambda(image=preprocessing_fn),\n        albu.Lambda(image=to_tensor, mask=to_tensor),\n    ]\n    return albu.Compose(_transform)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss"},{"metadata":{"trusted":true},"cell_type":"code","source":"class FirstHeadDiceSecondHeadBCE(smp.utils.base.Loss):\n    def __init__(self, eps=1., beta=1., activation=None, ignore_channels=None, **kwargs):\n        super().__init__(**kwargs)\n        self.dice_loss = smp.utils.losses.DiceLoss(eps=1., beta=1., activation=None, ignore_channels=None, **kwargs)\n        self.bce = smp.utils.losses.BCEWithLogitsLoss()\n        \n    def forward(self, y_pr, y_gt):\n        return self.dice_loss(y_pr[0], y_gt[0]) + self.bce(y_pr[1], y_gt[1])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyEpoch(smp.utils.train.Epoch):\n    def run(self, dataloader):\n\n        self.on_epoch_start()\n\n        logs = {}\n        loss_meter = smp.utils.meter.AverageValueMeter()\n        metrics_meters = {metric.__name__: smp.utils.meter.AverageValueMeter() for metric in self.metrics}\n\n        with tqdm(dataloader, desc=self.stage_name, file=sys.stdout, disable=not (self.verbose)) as iterator:\n            for x, y in iterator:\n                x, y = x.to(self.device), (y[0].to(self.device), y[1].to(self.device))\n                loss, y_pred = self.batch_update(x, y)\n\n                # update loss logs\n                loss_value = loss.cpu().detach().numpy()\n                loss_meter.add(loss_value)\n                loss_logs = {self.loss.__name__: loss_meter.mean}\n                logs.update(loss_logs)\n\n                # update metrics logs\n                for metric_fn in self.metrics:\n                    metric_value = metric_fn(y_pred[0], y[0]).cpu().detach().numpy()\n                    metrics_meters[metric_fn.__name__].add(metric_value)\n                metrics_logs = {k: v.mean for k, v in metrics_meters.items()}\n                logs.update(metrics_logs)\n\n                if self.verbose:\n                    s = self._format_logs(logs)\n                    iterator.set_postfix_str(s)\n\n        return logs\n    \nclass TrainEpoch(MyEpoch):\n\n    def __init__(self, model, loss, metrics, optimizer, device='cpu', verbose=True):\n        super().__init__(\n            model=model,\n            loss=loss,\n            metrics=metrics,\n            stage_name='train',\n            device=device,\n            verbose=verbose,\n        )\n        self.optimizer = optimizer\n\n    def on_epoch_start(self):\n        self.model.train()\n\n    def batch_update(self, x, y):\n        self.optimizer.zero_grad()\n        prediction = self.model.forward(x)\n        loss = self.loss(prediction, y)\n        loss.backward()\n        self.optimizer.step()\n        return loss, prediction\n    \nclass ValidEpoch(MyEpoch):\n\n    def __init__(self, model, loss, metrics, device='cpu', verbose=True):\n        super().__init__(\n            model=model,\n            loss=loss,\n            metrics=metrics,\n            stage_name='valid',\n            device=device,\n            verbose=verbose,\n        )\n\n    def on_epoch_start(self):\n        self.model.eval()\n\n    def batch_update(self, x, y):\n        with torch.no_grad():\n            prediction = self.model.forward(x)\n            loss = self.loss(prediction, y)\n        return loss, prediction","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"ENCODER = 'mobilenet_v2'\nENCODER_WEIGHTS = 'imagenet'\nDEVICE = 'cuda'\n\nACTIVATION = 'sigmoid'\n\naux_params=dict(\n    pooling='avg',             \n    dropout=0.2,               \n    activation=None,      \n    classes=NUM_CLASSES,                 \n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = smp.Unet(\n    encoder_name=ENCODER, \n    encoder_weights=ENCODER_WEIGHTS, \n    classes=NUM_CLASSES, \n    activation=ACTIVATION,\n    aux_params=aux_params\n)\n\npreprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = UnetDataset(\n    root_path_train,\n    train_df,\n    IMG_HEIGHT,\n    IMG_WIDTH, \n    preprocessing=get_preprocessing(preprocessing_fn),\n    augmentation=get_training_augmentation()\n)\n\nvalid_dataset = UnetDataset(\n    root_path_train,\n    val_df,\n    IMG_HEIGHT,\n    IMG_WIDTH, \n    preprocessing=get_preprocessing(preprocessing_fn),\n)\n\n\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=N_WORKERS, collate_fn=collate_function)\nvalid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=N_WORKERS, collate_fn=collate_function)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"loss = FirstHeadDiceSecondHeadBCE()\nmetrics = [\n    smp.utils.metrics.IoU(threshold=0.5),\n]\n\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# create epoch runners \n# it is a simple loop of iterating over dataloader`s samples\ntrain_epoch = TrainEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    optimizer=optimizer,\n    device=DEVICE,\n    verbose=True,\n)\n\nvalid_epoch = ValidEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    device=DEVICE,\n    verbose=True,\n)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model.state_dict(), 'best_model.pth')\nmax_score = 0\n\nfor i in range(0, 2):\n    \n    print('\\nEpoch: {}'.format(i))\n    train_logs = train_epoch.run(train_loader)\n    valid_logs = valid_epoch.run(valid_loader)\n    \n    # do something (save model, change lr, etc.)\n    if max_score < valid_logs['iou_score']:\n        max_score = valid_logs['iou_score']\n        torch.save(model.state_dict(), 'best_model.pth')\n        print('Model saved!')","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.6.6"}},"nbformat":4,"nbformat_minor":1}