{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Pip installs"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install segmentation-models-pytorch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Imports"},{"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 collections\nimport segmentation_models_pytorch as smp\nimport albumentations as albu\n\nfrom torch.utils.data import DataLoader\nfrom sklearn.model_selection import train_test_split\nfrom glob import glob\nfrom os import path\nfrom PIL import Image\nfrom tqdm import tqdm\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":"ROOT_PATH_TRAIN = '/kaggle/input/imaterialist-fashion-2019-FGVC6/train'\nDF_PATH_TRAIN = '/kaggle/input/imaterialist-fashion-2019-FGVC6/train.csv'\nPATH_TO_MODEL_WEIGHTS = '/kaggle/input/za-cho-takoe-testovoe/best_model.pth'\n\nIMAGE_PREDICTION_SIZE = (256, 256)\nN_CLASSES = 46","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Utils fucntions"},{"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":"# Train/Test split"},{"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":"class 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        \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            \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\n\n    def __len__(self):\n        return len(self.image_info)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Augmentations"},{"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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Preprocessing"},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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":"# Model"},{"metadata":{"trusted":true},"cell_type":"code","source":"ENCODER = 'mobilenet_v2'\nENCODER_WEIGHTS = 'imagenet'\nDEVICE = 'cuda'\n\nACTIVATION = 'sigmoid'","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=N_CLASSES, \n    activation=ACTIVATION,\n)\n\npreprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"train_dataset = UnetDataset(\n    ROOT_PATH_TRAIN,\n    train_df,\n    IMAGE_PREDICTION_SIZE[0],\n    IMAGE_PREDICTION_SIZE[1], \n    preprocessing=get_preprocessing(preprocessing_fn),\n    augmentation=get_training_augmentation()\n)\n\nvalid_dataset = UnetDataset(\n    ROOT_PATH_TRAIN,\n    val_df,\n    IMAGE_PREDICTION_SIZE[0],\n    IMAGE_PREDICTION_SIZE[1], \n    preprocessing=get_preprocessing(preprocessing_fn),\n)\n\n\ntrain_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=2)\nvalid_loader = DataLoader(valid_dataset, batch_size=32, shuffle=False, num_workers=2)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss Metric Optimizer"},{"metadata":{"trusted":true},"cell_type":"code","source":"loss = smp.utils.losses.DiceLoss()\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":{},"cell_type":"markdown","source":"# Epoch class"},{"metadata":{"trusted":true},"cell_type":"code","source":"# create epoch runners \n# it is a simple loop of iterating over dataloader`s samples\ntrain_epoch = smp.utils.train.TrainEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    optimizer=optimizer,\n    device=DEVICE,\n    verbose=True,\n)\n\nvalid_epoch = smp.utils.train.ValidEpoch(\n    model, \n    loss=loss, \n    metrics=metrics, \n    device=DEVICE,\n    verbose=True,\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"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":{"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}