{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Imports"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install segmentation-models-pytorch","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","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 albumentations as albu\nimport torchvision\nimport segmentation_models_pytorch as smp\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\nfrom torchvision.models.detection.faster_rcnn import FastRCNNPredictor\nfrom torchvision.models.detection.mask_rcnn import MaskRCNNPredictor\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 = 16\nN_WORKERS = 2\n\nDEVICE = 'cuda'\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":{"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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","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\n\n\n\ndef get_unique_class_id_df(inital_df):\n    temp_df = inital_df.groupby(['ImageId','ClassId'])['EncodedPixels'].agg(lambda x: ' '.join(list(x))).reset_index()\n    size_df = inital_df.groupby(['ImageId','ClassId'])['Height', 'Width'].mean().reset_index()\n    temp_df = temp_df.merge(size_df, on=['ImageId','ClassId'], how='left')\n    \n    return temp_df","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":{"trusted":true},"cell_type":"code","source":"train_df, val_df = get_unique_class_id_df(train_df), get_unique_class_id_df(val_df)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"class FashionDataset(torch.utils.data.Dataset):\n    def __init__(self, image_dir, df, height, width, transforms=None):\n        self.transforms = transforms\n        self.image_dir = image_dir\n        self.df = df\n        self.height = height\n        self.width = width\n        self.image_info = collections.defaultdict(dict)\n        self.df['CategoryId'] = self.df.ClassId.apply(lambda x: str(x).split(\"_\")[0])\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        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        self.img2tensor = torchvision.transforms.ToTensor()\n\n    def __getitem__(self, idx):\n        # load images ad masks\n        img_path = self.image_info[idx][\"image_path\"]\n        img = Image.open(img_path).convert(\"RGB\")\n        img = img.resize((self.width, self.height), resample=Image.BILINEAR)\n\n        info = self.image_info[idx]\n        mask = np.zeros((len(info['annotations']), self.width, self.height), dtype=np.uint8)\n        labels = []\n        for m, (annotation, label) in enumerate(zip(info['annotations'], info['labels'])):\n            sub_mask = rle_decode(annotation, (info['orig_height'], info['orig_width']))\n            sub_mask = Image.fromarray(sub_mask)\n            sub_mask = sub_mask.resize((self.width, self.height), resample=Image.BILINEAR)\n            mask[m, :, :] = sub_mask\n            labels.append(int(label) + 1)\n\n        num_objs = len(labels)\n        boxes = []\n        new_labels = []\n        new_masks = []\n\n        for i in range(num_objs):\n            try:\n                pos = np.where(mask[i, :, :])\n                xmin = np.min(pos[1])\n                xmax = np.max(pos[1])\n                ymin = np.min(pos[0])\n                ymax = np.max(pos[0])\n                if abs(xmax - xmin) >= 20 and abs(ymax - ymin) >= 20:\n                    boxes.append([xmin, ymin, xmax, ymax])\n                    new_labels.append(labels[i])\n                    new_masks.append(mask[i, :, :])\n            except ValueError:\n                continue\n\n        if len(new_labels) == 0:\n            boxes.append([0, 0, 20, 20])\n            new_labels.append(0)\n            new_masks.append(mask[0, :, :])\n\n        nmx = np.zeros((len(new_masks), self.width, self.height), dtype=np.uint8)\n        for i, n in enumerate(new_masks):\n            nmx[i, :, :] = n\n\n        target = {}\n        target[\"boxes\"] = torch.as_tensor(boxes, dtype=torch.float32)\n        target[\"labels\"] = torch.as_tensor(new_labels, dtype=torch.int64)\n        target[\"masks\"] = torch.as_tensor(nmx, dtype=torch.uint8)\n        \n        img = self.img2tensor(img)\n\n        if self.transforms is not None:\n            img, target = self.transforms(img, target)\n\n        return img, target\n\n    def __len__(self):\n        return len(self.image_info)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def custom_collate(batch):\n    images = []\n    labels = []\n    for img, label in batch:\n        images.append(img)\n        labels.append(label)\n        \n    return images, labels","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"datset = FashionDataset(image_dir=root_path_train, \n                        df=train_df, \n                        height=IMG_HEIGHT, \n                        width=IMG_WIDTH)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data_loader = torch.utils.data.DataLoader(\n    datset, batch_size=4, shuffle=True, num_workers=2,\n    collate_fn=custom_collate)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Train epoch"},{"metadata":{"trusted":true},"cell_type":"code","source":"class MyEpoch(smp.utils.train.Epoch):\n    def _to_device(self):\n        self.model.to(self.device)\n        \n    def run(self, dataloader):\n\n        self.on_epoch_start()\n\n        logs = {}\n        loss_meter = smp.utils.meter.AverageValueMeter()\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 = list(map(lambda x_el: x_el.to(self.device), x))\n                y = list(map(lambda y_el: {k:v.to(self.device) for k,v in y_el.items()}, y))\n                loss = 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 = {'loss': loss_meter.mean}\n                logs.update(loss_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        loss = self.model(x, y)\n        loss = sum(l for l in loss.values())\n        loss.backward()\n        self.optimizer.step()\n        return loss\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"num_classes = NUM_CLASSES + 1\ndevice = torch.device(DEVICE)\n\nmodel_ft = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)\nin_features = model_ft.roi_heads.box_predictor.cls_score.in_features\nmodel_ft.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)\nin_features_mask = model_ft.roi_heads.mask_predictor.conv5_mask.in_channels\nhidden_layer = 256\nmodel_ft.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"for param in model_ft.parameters():\n    param.requires_grad = True","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Test it"},{"metadata":{"trusted":true},"cell_type":"code","source":"optimizer = torch.optim.Adam(model_ft.parameters(), lr=0.001)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_epoch = TrainEpoch(\n    model_ft, \n    loss=None, \n    metrics=None, \n    optimizer=optimizer,\n    device=DEVICE,\n    verbose=True,\n)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model_ft.state_dict(), 'best_model.pth')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_epoch.run(data_loader)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.save(model_ft.state_dict(), 'best_model.pth')","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}