{"cells":[{"metadata":{"id":"1_dPgKgwSQfp","_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!pip install albumentations==0.4.6\n!pip install git+https://github.com/qubvel/segmentation_models.pytorch","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Importing libraries"},{"metadata":{"id":"bK6emFemSkOo","_kg_hide-input":true,"trusted":true},"cell_type":"code","source":"import os\nimport time\nimport io\nimport base64\n\nimport numpy as np\nimport pandas as pd\nimport cv2\n\nimport matplotlib.pyplot as plt\nimport matplotlib.colors as mpl_colors\nimport seaborn as sns\n\nfrom sklearn.model_selection import train_test_split\n\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nfrom albumentations.pytorch import ToTensor, ToTensorV2 \nfrom albumentations import (HorizontalFlip,\n                            VerticalFlip,\n                            Normalize,\n                            Compose)\n\nfrom segmentation_models_pytorch.unet import Unet\n\nfrom IPython.display import HTML\nfrom IPython.display import clear_output\nimport warnings\nwarnings.simplefilter(\"ignore\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### Helper functions"},{"metadata":{"trusted":true,"_kg_hide-input":true},"cell_type":"code","source":"# utils\ndef get_one_slice_data(img_name: str,\n                       mask_name: str,\n                       root_imgs_path: str = \"images/\",\n                       root_masks_path: str = \"masks/\",) -> np.ndarray:\n\n    img_path = os.path.join(root_imgs_path, img_name)\n    mask_path = os.path.join(root_masks_path, mask_name)\n    one_slice_img = cv2.imread(img_path)#[:,:,0] uncomment for grayscale\n    one_slice_mask = cv2.imread(mask_path)\n    one_slice_mask[one_slice_mask < 240] = 0  # remove artifacts\n    one_slice_mask[one_slice_mask >= 240] = 255\n\n    return one_slice_img, one_slice_mask\n\n\ndef get_id_predictions(net: nn.Module,\n                       ct_scan_id_df: pd.DataFrame,\n                       root_imgs_dir: str,\n                       treshold: float = 0.3) -> list:\n\n    \"\"\"\n    Factory for getting predictions and storing them and images in lists as uint8 images.\n    Params:\n        net: model for prediction.\n        ct_scan_id_df: df with unique patient id.\n        root_imgs_dir: root path for images.\n        treshold: threshold for probabilities.\n    \"\"\"\n    sigmoid = lambda x: 1 / (1 + np.exp(-x))\n    images = []\n    predictions = []\n    net.eval()\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    print(\"device:\", device)\n    with torch.no_grad():\n        for idx in range(len(ct_scan_id_df)):\n            img_name = ct_scan_id_df.loc[idx, \"ImageId\"]\n            path = os.path.join(root_imgs_dir, img_name)\n\n            img_ = cv2.imread(path)\n    \n            img = Normalize().apply(img_)\n            tensor = torch.FloatTensor(img).permute(2, 0, 1).unsqueeze(0)\n            prediction = net.forward(tensor.to(device))\n            prediction = prediction.cpu().detach().numpy()\n            prediction = prediction.squeeze(0).transpose(1, 2, 0)\n            prediction = sigmoid(prediction)\n            prediction = (prediction >= treshold).astype(np.float32)\n\n            predictions.append((prediction * 255).astype(\"uint8\"))\n            images.append(img_)\n\n    return images, predictions\n\n\n# Save image in original resolution\n# helpful link - https://stackoverflow.com/questions/34768717/matplotlib-unable-to-save-image-in-same-resolution-as-original-image\n\ndef get_overlaid_masks_on_image(\n                one_slice_image: np.ndarray,\n                one_slice_mask: np.ndarray, \n                w: float = 512,\n                h: float = 512, \n                dpi: float = 100,\n                write: bool = False,\n                path_to_save: str = '/content/',\n                name_to_save: str = 'img_name'):\n    \"\"\"overlap masks on image and save this as a new image.\"\"\"\n\n    path_to_save_ = os.path.join(path_to_save, name_to_save)\n    lung, heart, trachea = [one_slice_mask[:, :, i] for i in range(3)]\n    figsize = (w / dpi), (h / dpi)\n    fig = plt.figure(figsize=(figsize))\n    fig.add_axes([0, 0, 1, 1])\n \n    # image\n    plt.imshow(one_slice_image, cmap=\"bone\")\n\n    # overlaying segmentation masks\n    plt.imshow(np.ma.masked_where(lung == False, lung),\n            cmap='cool', alpha=0.3)\n    plt.imshow(np.ma.masked_where(heart == False, heart),\n            cmap='autumn', alpha=0.3)\n    plt.imshow(np.ma.masked_where(trachea == False, trachea),\n               cmap='autumn_r', alpha=0.3) \n\n    #plt.axis(\"off\")                                            ## dont work in kaggle kernel\n    fig.axes[0].get_xaxis().set_visible(False)                  ## work in kaggle kernel\n    fig.axes[0].get_yaxis().set_visible(False)\n    \n\n    fig.savefig(f\"{path_to_save_}.png\",bbox_inches='tight', \n                pad_inches=0, dpi=dpi,  format=\"png\")\n    \n    if write:\n        plt.close()\n    else:\n        plt.show()\n        \ndef get_overlaid_masks_on_full_ctscan(ct_scan_id_df: pd.DataFrame, \n                                      path_to_save: str,\n                                      root_imgs_dir: str,\n                                      root_masks_dir: str,\n                                     ):\n    \"\"\"\n    Creating images with overlaid masks on each slice of CT scan.\n    Params:\n         ct_scan_id_df: df with unique patient id.\n         path_to_save: path to save images.\n    \"\"\"\n    num_slice = len(ct_scan_id_df)\n    for slice_ in range(num_slice):\n        img_name = ct_scan_id_df.loc[slice_, \"ImageId\"]\n        mask_name = ct_scan_id_df.loc[slice_, \"MaskId\"]\n        one_slice_img, one_slice_mask = get_one_slice_data(img_name, mask_name,\n                                                           root_imgs_dir, root_masks_dir)\n        get_overlaid_masks_on_image(one_slice_img,\n                                one_slice_mask,\n                                write=True, \n                                path_to_save=path_to_save,\n                                name_to_save=str(slice_)\n                                )\n\ndef create_video(path_to_imgs: str, video_name: str, framerate: int):\n    \"\"\"\n    Create video from images.\n    Params:\n        path_to_imgs: path to dir with images.\n        video_name: name for saving video.\n        framerate: num frames per sec in video.\n    \"\"\"\n    img_names = sorted(os.listdir(path_to_imgs), key=lambda x: int(x[:-4]))  # img_name must be numbers\n    img_path = os.path.join(path_to_imgs, img_names[0])\n    frame_width, frame_height, _ = cv2.imread(img_path).shape\n    fourc = cv2.VideoWriter_fourcc(*'XVID')           ## MP4V - dont work in kaggle kernel\n    video = cv2.VideoWriter(video_name + \".avi\",      ##'mp4' \n                            fourc, \n                            framerate, \n                            (frame_width, frame_height))\n\n    for img_name in img_names:\n        img_path = os.path.join(path_to_imgs, img_name)\n        image = cv2.imread(img_path)\n        video.write(image)\n            \n    cv2.destroyAllWindows()\n    video.release()\n\n    \ndef compute_scores_per_classes(model,\n                               dataloader,\n                               classes):\n    \"\"\"\n    Compute Dice and Jaccard coefficients for each class.\n    Params:\n        model: neural net for make predictions.\n        dataloader: dataset object to load data from.\n        classes: list with classes.\n        Returns: dictionaries with dice and jaccard coefficients for each class for each slice.\n    \"\"\"\n    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n    dice_scores_per_classes = {key: list() for key in classes}\n    iou_scores_per_classes = {key: list() for key in classes}\n\n    with torch.no_grad():\n        for i, (imgs, targets) in enumerate(dataloader):\n            imgs, targets = imgs.to(device), targets.to(device)\n            logits = model(imgs)\n            logits = logits.detach().cpu().numpy()\n            targets = targets.detach().cpu().numpy()\n            \n            dice_scores = dice_coef_metric_per_classes(logits, targets)\n            iou_scores = jaccard_coef_metric_per_classes(logits, targets)\n\n            for key in dice_scores.keys():\n                dice_scores_per_classes[key].extend(dice_scores[key])\n\n            for key in iou_scores.keys():\n                iou_scores_per_classes[key].extend(iou_scores[key])\n\n    return dice_scores_per_classes, iou_scores_per_classes\n\n\n# loss\ndef dice_coef_metric_per_classes(probabilities: np.ndarray,\n                                    truth: np.ndarray,\n                                    treshold: float = 0.5,\n                                    eps: float = 1e-9,\n                                    classes: list = ['lung', 'heart', 'trachea']) -> np.ndarray:\n    \"\"\"\n    Calculate Dice score for data batch and for each class.\n    Params:\n        probobilities: model outputs after activation function.\n        truth: model targets.\n        threshold: threshold for probabilities.\n        eps: additive to refine the estimate.\n        classes: list with name classes.\n        Returns: dict with dice scores for each class.\n    \"\"\"\n    scores = {key: list() for key in classes}\n    num = probabilities.shape[0]\n    num_classes = probabilities.shape[1]\n    predictions = (probabilities >= treshold).astype(np.float32)\n    assert(predictions.shape == truth.shape)\n\n    for i in range(num):\n        for class_ in range(num_classes):\n            prediction = predictions[i][class_]\n            truth_ = truth[i][class_]\n            intersection = 2.0 * (truth_ * prediction).sum()\n            union = truth_.sum() + prediction.sum()\n            if truth_.sum() == 0 and prediction.sum() == 0:\n                 scores[classes[class_]].append(1.0)\n            else:\n                scores[classes[class_]].append((intersection + eps) / union)\n                \n    return scores\n\n\ndef jaccard_coef_metric_per_classes(probabilities: np.ndarray,\n               truth: np.ndarray,\n               treshold: float = 0.5,\n               eps: float = 1e-9,\n               classes: list = ['lung', 'heart', 'trachea']) -> np.ndarray:\n    \"\"\"\n    Calculate Jaccard index for data batch and for each class.\n    Params:\n        probobilities: model outputs after activation function.\n        truth: model targets.\n        threshold: threshold for probabilities.\n        eps: additive to refine the estimate.\n        classes: list with name classes.\n        Returns: dict with jaccard scores for each class.\"\n    \"\"\"\n    scores = {key: list() for key in classes}\n    num = probabilities.shape[0]\n    num_classes = probabilities.shape[1]\n    predictions = (probabilities >= treshold).astype(np.float32)\n    assert(predictions.shape == truth.shape)\n\n    for i in range(num):\n        for class_ in range(num_classes):\n            prediction = predictions[i][class_]\n            truth_ = truth[i][class_]\n            intersection = (prediction * truth_).sum()\n            union = (prediction.sum() + truth_.sum()) - intersection + eps\n            if truth_.sum() == 0 and prediction.sum() == 0:\n                 scores[classes[class_]].append(1.0)\n            else:\n                scores[classes[class_]].append((intersection + eps) / union)\n\n    return scores\n\n\n# visualizer\ndef show_data_augmentations(img_tensor: torch.Tensor,\n                            mask_tensor: torch.Tensor,\n                            mean: tuple = (0.485, 0.456, 0.406),\n                            std: tuple = (0.229, 0.224, 0.225),\n                            labels: list=[\"image\", \"lung\", \"heart\", \"trachea\"]):\n    \n    img = img_tensor.numpy().transpose(1, 2, 0)\n    img = (img * std + mean).astype(\"float32\")\n    img = np.clip(img, 0, 1)\n    mask = mask_tensor.numpy().transpose(1, 2, 0)\n    data_to_plot = [img, *[mask[:,:, i] for i in range(3)]]\n\n    fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(15, 8))\n    for i, ax in enumerate(axes):\n        ax.imshow(data_to_plot[i])\n        ax.set_title(labels[i])\n\n    plt.show()\n    \n    \ndef show_video(video_path: str):\n    \"\"\"\n    show video in jupyter notebook, agent interaction in environment.\n    Takes - path to video file.\n    Returns - html video player in jupyter notebook.\n    \"\"\"  \n    video = io.open(video_path, 'r+b').read()\n    encoded = base64.b64encode(video)\n\n    return HTML(data='''<video alt=\"test\" controls>\n    <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" /> </video>'''\n    .format(encoded.decode('ascii')))\n\n\ndef get_color_info(classes: list = ['lung', 'heart', 'trachea']):\n\n    def get_color(cmap):\n        new_cmap = mpl_colors.LinearSegmentedColormap.from_list(\n            (cmap.name, 0.0, 0.0,),\n            cmap(np.linspace(0.0, 0.0, 100))\n            )\n        return new_cmap\n\n    colormaps = [plt.get_cmap(cmap_name) for cmap_name in \n                ['cool', 'autumn', 'autumn_r']\n                ]\n    arr = np.linspace(0, 50, 100).reshape((10, 10))\n    fig, ax = plt.subplots(1, 3, figsize=(15, 5))\n    for i in range(3):\n        cmap = get_color(colormaps[i])\n        ax[i].axis('off')\n        ax[i].imshow(arr, cmap=cmap)\n        ax[i].set_title(classes[i], fontsize=15)\n\n    fig.suptitle(\"Color definition\", fontsize=20, y=0.99)\n    fig.savefig(f\"color_definition.png\", bbox_inches='tight', \n                pad_inches=0.2, dpi=100, format=\"png\")\n    \n    plt.savefig(f\"color_definition.svg\", bbox_inches='tight',\n                pad_inches=0.2, dpi=100, format=\"svg\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"config with global variables"},{"metadata":{"id":"8WmHqD50SkRe","trusted":true},"cell_type":"code","source":"class GlobalConfig:\n    def __init__(self):\n        self.seed = 555\n        self.path_to_csv = '../input/chest-ct-segmentation/train.csv'\n        self.path_to_imgs_dir = '../input/chest-ct-segmentation/images/images'\n        self.path_to_masks_dir = '../input/chest-ct-segmentation/masks/masks'\n        self.pretrained_model_path = '../input/chest-ct-segmentation/pretrained_model/pretrained_model/model_100_epoch.pth'\n        self.train_logs_path = '../input/chest-ct-segmentation/pretrained_model/pretrained_model/train_log_100_epoch.csv'\n\n\ndef seed_everything(seed: int):\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    np.random.seed(seed)\n    \nconfig = GlobalConfig()\nseed_everything(config.seed)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset and Dataloader"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LungsDataset(Dataset):\n    def __init__(self, \n                 imgs_dir: str,\n                 masks_dir:str,\n                 df: pd.DataFrame,\n                 phase: str):\n        \"\"\"Initialization.\"\"\"\n        self.root_imgs_dir = imgs_dir\n        self.root_masks_dir = masks_dir\n        self.df = df\n        self.augmentations = get_augmentations(phase)\n    \n    def __len__(self):\n        return len(self.df)\n\n    def __getitem__(self, idx):\n        img_name = self.df.loc[idx, \"ImageId\"]\n        mask_name = self.df.loc[idx, \"MaskId\"]\n        img_path = os.path.join(self.root_imgs_dir, img_name)\n        mask_path = os.path.join(self.root_masks_dir, mask_name)\n        img = cv2.imread(img_path)\n        mask = cv2.imread(mask_path)\n        mask[mask < 240] = 0    # remove artifacts\n        mask[mask > 0] = 1\n\n        augmented = self.augmentations(image=img, \n                                       mask=mask.astype(np.float32))\n        img = augmented['image']\n        mask = augmented['mask'].permute(2, 0, 1)\n\n        return img, mask\n\n\ndef get_augmentations(phase,\n                   mean: tuple = (0.485, 0.456, 0.406),\n                   std: tuple = (0.229, 0.224, 0.225),):\n    list_transforms = []\n    if phase == \"train\":\n        list_transforms.extend(\n            [\n                VerticalFlip(p=0.5), \n            ]\n        )\n    list_transforms.extend(\n        [\n            Normalize(mean=mean, std=std, p=1),\n            #ToTensor(num_classes=3, sigmoid=False),\n            ToTensorV2(),\n        ]\n    )\n    list_trfms = Compose(list_transforms)\n    return list_trfms\n\n\ndef get_dataloader(\n    imgs_dir: str,\n    masks_dir: str,\n    path_to_csv: str,\n    phase: str,\n    batch_size: int = 8,\n    num_workers: int = 6,\n    test_size: float = 0.2,\n):\n    '''Returns: dataloader for the model training'''\n    df = pd.read_csv(path_to_csv)\n    \n\n    train_df, val_df = train_test_split(df, \n                                          test_size=test_size, \n                                          random_state=69)\n    train_df, val_df = train_df.reset_index(drop=True), val_df.reset_index(drop=True)\n\n    df = train_df if phase == \"train\" else val_df\n    image_dataset = LungsDataset(imgs_dir, masks_dir, df, phase)\n    dataloader = DataLoader(\n        image_dataset,\n        batch_size=batch_size,\n        num_workers=num_workers,\n        pin_memory=True,\n        shuffle=True,   \n    )\n\n    return dataloader","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Loss and Metrics"},{"metadata":{"trusted":true},"cell_type":"code","source":"def dice_coef_metric(probabilities: torch.Tensor,\n                     truth: torch.Tensor,\n                     treshold: float = 0.5,\n                     eps: float = 1e-9) -> np.ndarray:\n    \"\"\"\n    Calculate Dice score for data batch.\n    Params:\n        probobilities: model outputs after activation function.\n        truth: truth values.\n        threshold: threshold for probabilities.\n        eps: additive to refine the estimate.\n        Returns: dice score aka f1.\n    \"\"\"\n    scores = []\n    num = probabilities.shape[0]\n    predictions = (probabilities >= treshold).float()\n    assert(predictions.shape == truth.shape)\n    for i in range(num):\n        prediction = predictions[i]\n        truth_ = truth[i]\n        intersection = 2.0 * (truth_ * prediction).sum()\n        union = truth_.sum() + prediction.sum()\n        if truth_.sum() == 0 and prediction.sum() == 0:\n            scores.append(1.0)\n        else:\n            scores.append((intersection + eps) / union)\n    return np.mean(scores)\n\n\ndef jaccard_coef_metric(probabilities: torch.Tensor,\n               truth: torch.Tensor,\n               treshold: float = 0.5,\n               eps: float = 1e-9) -> np.ndarray:\n    \"\"\"\n    Calculate Jaccard index for data batch.\n    Params:\n        probobilities: model outputs after activation function.\n        truth: truth values.\n        threshold: threshold for probabilities.\n        eps: additive to refine the estimate.\n        Returns: jaccard score aka iou.\"\n    \"\"\"\n    scores = []\n    num = probabilities.shape[0]\n    predictions = (probabilities >= treshold).float()\n    assert(predictions.shape == truth.shape)\n\n    for i in range(num):\n        prediction = predictions[i]\n        truth_ = truth[i]\n        intersection = (prediction * truth_).sum()\n        union = (prediction.sum() + truth_.sum()) - intersection + eps\n        if truth_.sum() == 0 and prediction.sum() == 0:\n            scores.append(1.0)\n        else:\n            scores.append((intersection + eps) / union)\n    return np.mean(scores)\n\n\nclass Meter:\n    '''factory for storing and updating iou and dice scores.'''\n    def __init__(self, treshold: float = 0.5):\n        self.threshold: float = treshold\n        self.dice_scores: list = []\n        self.iou_scores: list = []\n    \n    def update(self, logits: torch.Tensor, targets: torch.Tensor):\n        \"\"\"\n        Takes: logits from output model and targets,\n        calculates dice and iou scores, and stores them in lists.\n        \"\"\"\n        probs = torch.sigmoid(logits)\n        dice = dice_coef_metric(probs, targets, self.threshold)\n        iou = jaccard_coef_metric(probs, targets, self.threshold)\n        \n        self.dice_scores.append(dice)\n        self.iou_scores.append(iou)\n    \n    def get_metrics(self) -> np.ndarray:\n        \"\"\"\n        Returns: the average of the accumulated dice and iou scores.\n        \"\"\"\n        dice = np.mean(self.dice_scores)\n        iou = np.mean(self.iou_scores)\n        return dice, iou\n    \n\nclass DiceLoss(nn.Module):\n    \"\"\"Calculate dice loss.\"\"\"\n    def __init__(self, eps: float = 1e-9):\n        super(DiceLoss, self).__init__()\n        self.eps = eps\n        \n    def forward(self,\n                logits: torch.Tensor,\n                targets: torch.Tensor) -> torch.Tensor:\n        \n        num = targets.size(0)\n        probability = torch.sigmoid(logits)\n        probability = probability.view(num, -1)\n        targets = targets.view(num, -1)\n        assert(probability.shape == targets.shape)\n        \n        intersection = 2.0 * (probability * targets).sum()\n        union = probability.sum() + targets.sum()\n        dice_score = (intersection + self.eps) / union\n        #print(\"intersection\", intersection, union, dice_score)\n        return 1.0 - dice_score\n        \n        \nclass BCEDiceLoss(nn.Module):\n    \"\"\"Compute objective loss: BCE loss + DICE loss.\"\"\"\n    def __init__(self):\n        super(BCEDiceLoss, self).__init__()\n        self.bce = nn.BCEWithLogitsLoss()\n        self.dice = DiceLoss()\n        \n    def forward(self, \n                logits: torch.Tensor,\n                targets: torch.Tensor) -> torch.Tensor:\n        assert(logits.shape == targets.shape)\n        dice_loss = self.dice(logits, targets)\n        bce_loss = self.bce(logits, targets)\n        \n        return bce_loss + dice_loss\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"metadata":{"id":"UWec2CrXUPjC","_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"model = Unet('efficientnet-b2', encoder_weights=\"imagenet\", classes=3, activation=None)","execution_count":null,"outputs":[]},{"metadata":{"id":"PWUHmzvbaKo0"},"cell_type":"markdown","source":"# Train Process"},{"metadata":{"trusted":true},"cell_type":"code","source":"class Trainer:\n    \"\"\"\n    Factory for training proccess.\n    Args:\n        display_plot: if True - plot train history after each epoch.\n        net: neural network for mask prediction.\n        criterion: factory for calculating objective loss.\n        optimizer: optimizer for weights updating.\n        phases: list with train and validation phases.\n        dataloaders: dict with data loaders for train and val phases.\n        imgs_dir: path to folder with images.\n        masks_dir: path to folder with imasks.\n        path_to_csv: path to csv file.\n        meter: factory for storing and updating metrics.\n        batch_size: data batch size for one step weights updating.\n        num_epochs: num weights updation for all data.\n        accumulation_steps: the number of steps after which the optimization step can be taken\n                    (https://www.kaggle.com/c/understanding_cloud_organization/discussion/105614).\n        lr: learning rate for optimizer.\n        scheduler: scheduler for control learning rate.\n        losses: dict for storing lists with losses for each phase.\n        jaccard_scores: dict for storing lists with jaccard scores for each phase.\n        dice_scores: dict for storing lists with dice scores for each phase.\n    \"\"\"\n    def __init__(self,\n                 net: nn.Module,\n                 criterion: nn.Module,\n                 lr: float,\n                 accumulation_steps: int,\n                 batch_size: int,\n                 num_epochs: int,\n                 imgs_dir: str,\n                 masks_dir: str,\n                 path_to_csv: str,\n                 display_plot: bool = True\n                ):\n\n        \"\"\"Initialization.\"\"\"\n        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'\n        print(\"device:\", self.device)\n        self.display_plot = display_plot\n        self.net = net\n        self.net = self.net.to(self.device)\n        self.criterion = criterion\n        self.optimizer = Adam(self.net.parameters(), lr=lr)\n        self.scheduler = ReduceLROnPlateau(self.optimizer, mode=\"min\",\n                                           patience=3, verbose=True)\n        self.accumulation_steps = accumulation_steps // batch_size\n        self.phases = [\"train\", \"val\"]\n        self.num_epochs = num_epochs\n\n        self.dataloaders = {\n            phase: get_dataloader(\n                imgs_dir = imgs_dir,\n                masks_dir = masks_dir,\n                path_to_csv = path_to_csv,\n                phase = phase,\n                batch_size = 8,\n                num_workers = 6\n            )\n            for phase in self.phases\n        }\n        self.best_loss = float(\"inf\")\n        self.losses = {phase: [] for phase in self.phases}\n        self.dice_scores = {phase: [] for phase in self.phases}\n        self.jaccard_scores = {phase: [] for phase in self.phases}\n         \n    def _compute_loss_and_outputs(self,\n                                  images: torch.Tensor,\n                                  targets: torch.Tensor):\n        images = images.to(self.device)\n        targets = targets.to(self.device)\n        logits = self.net(images)\n        loss = self.criterion(logits, targets)\n        return loss, logits\n        \n    def _do_epoch(self, epoch: int, phase: str):\n        print(f\"{phase} epoch: {epoch} | time: {time.strftime('%H:%M:%S')}\")\n\n        self.net.train() if phase == \"train\" else self.net.eval()\n        meter = Meter()\n        dataloader = self.dataloaders[phase]\n        total_batches = len(dataloader)\n        running_loss = 0.0\n        self.optimizer.zero_grad()\n        for itr, (images, targets) in enumerate(dataloader):\n            loss, logits = self._compute_loss_and_outputs(images, targets)\n            loss = loss / self.accumulation_steps\n            if phase == \"train\":\n                loss.backward()\n                if (itr + 1) % self.accumulation_steps == 0:\n                    self.optimizer.step()\n                    self.optimizer.zero_grad()\n            running_loss += loss.item()\n            meter.update(logits.detach().cpu(),\n                         targets.detach().cpu()\n                        )\n            \n        epoch_loss = (running_loss * self.accumulation_steps) / total_batches\n        epoch_dice, epoch_iou = meter.get_metrics()\n        \n        self.losses[phase].append(epoch_loss)\n        self.dice_scores[phase].append(epoch_dice)\n        self.jaccard_scores[phase].append(epoch_iou)\n\n        return epoch_loss\n        \n    def train(self):\n        for epoch in range(self.num_epochs):\n            self._do_epoch(epoch, \"train\")\n            with torch.no_grad():\n                val_loss = self._do_epoch(epoch, \"val\")\n                self.scheduler.step(val_loss)\n            if self.display_plot:\n                self._plot_train_history()\n                \n            if val_loss < self.best_loss:\n                print(f\"\\n{'#'*20}\\nSaved new checkpoint\\n{'#'*20}\\n\")\n                self.best_loss = val_loss\n                torch.save(self.net.state_dict(), \"best_model.pth\")\n            print()\n        self._save_train_history()\n            \n    def _plot_train_history(self):\n        data = [self.losses, self.dice_scores, self.jaccard_scores]\n        colors = ['deepskyblue', \"crimson\"]\n        labels = [\n            f\"\"\"\n            train loss {self.losses['train'][-1]}\n            val loss {self.losses['val'][-1]}\n            \"\"\",\n            \n            f\"\"\"\n            train dice score {self.dice_scores['train'][-1]}\n            val dice score {self.dice_scores['val'][-1]} \n            \"\"\", \n                  \n            f\"\"\"\n            train jaccard score {self.jaccard_scores['train'][-1]}\n            val jaccard score {self.jaccard_scores['val'][-1]}\n            \"\"\",\n        ]\n        \n        clear_output(True)\n        with plt.style.context(\"seaborn-dark-palette\"):\n            fig, axes = plt.subplots(3, 1, figsize=(8, 10))\n            for i, ax in enumerate(axes):\n                ax.plot(data[i]['val'], c=colors[0], label=\"val\")\n                ax.plot(data[i]['train'], c=colors[-1], label=\"train\")\n                ax.set_title(labels[i])\n                ax.legend(loc=\"upper right\")\n                \n            plt.tight_layout()\n            plt.show()\n            \n    def load_predtrain_model(self,\n                             state_path: str):\n        self.net.load_state_dict(torch.load(state_path))\n        print(\"Predtrain model loaded\")\n        \n    def _save_train_history(self):\n        \"\"\"writing model weights and training logs to files.\"\"\"\n        torch.save(self.net.state_dict(),\n                   f\"last_epoch_model.pth\")\n\n        logs_ = [self.losses, self.dice_scores, self.jaccard_scores]\n        log_names_ = [\"_loss\", \"_dice\", \"_jaccard\"]\n        logs = [logs_[i][key] for i in list(range(len(logs_)))\n                         for key in logs_[i]]\n        log_names = [key+log_names_[i] \n                     for i in list(range(len(logs_))) \n                     for key in logs_[i]\n                    ]\n        pd.DataFrame(\n            dict(zip(log_names, logs))\n        ).to_csv(\"train_log.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{"id":"ULnA0lkgUPmI","outputId":"d83b2393-ca73-40dc-cd12-ab02b735e240","trusted":true},"cell_type":"code","source":"trainer = Trainer(net=model,\n                  criterion=BCEDiceLoss(),\n                  lr=8e-5,\n                  accumulation_steps=32,\n                  batch_size=8,\n                  num_epochs=1,\n                  imgs_dir = config.path_to_imgs_dir,\n                  masks_dir = config.path_to_masks_dir,\n                  path_to_csv = config.path_to_csv,)\n\nif config.pretrained_model_path is not None:\n    trainer.load_predtrain_model(config.pretrained_model_path)\n    \n    # if need - load the logs.      \n    train_logs = pd.read_csv(config.train_logs_path)\n    trainer.losses[\"train\"] =  train_logs.loc[:, \"train_loss\"].to_list()\n    trainer.losses[\"val\"] =  train_logs.loc[:, \"val_loss\"].to_list()\n    trainer.dice_scores[\"train\"] = train_logs.loc[:, \"train_dice\"].to_list()\n    trainer.dice_scores[\"val\"] = train_logs.loc[:, \"val_dice\"].to_list()\n    trainer.jaccard_scores[\"train\"] = train_logs.loc[:, \"train_jaccard\"].to_list()\n    trainer.jaccard_scores[\"val\"] = train_logs.loc[:, \"val_jaccard\"].to_list()","execution_count":null,"outputs":[]},{"metadata":{"id":"iM-3o86EUPre","outputId":"34dd86f5-6f20-4a41-ae37-e17954ce21d4","trusted":true},"cell_type":"code","source":"%%time\ntrainer.train()","execution_count":null,"outputs":[]},{"metadata":{"id":"okfuiTuAaGLp"},"cell_type":"markdown","source":"# Experiments and Results"},{"metadata":{"id":"H0AgOf_SbgYF","trusted":true},"cell_type":"code","source":"val_dataloader = get_dataloader(\n    imgs_dir=config.path_to_imgs_dir,\n    masks_dir=config.path_to_masks_dir,\n    path_to_csv=config.path_to_csv,\n    phase = \"val\",\n    batch_size = 8,\n    num_workers = 6,\n    test_size = 0.2,\n)","execution_count":null,"outputs":[]},{"metadata":{"id":"TG9fX_mzaNS5","outputId":"438b6e3d-5172-4d08-fa65-cc19b955268a","trusted":true},"cell_type":"code","source":"%%time\ndice_scores_per_classes, iou_scores_per_classes = compute_scores_per_classes(\n    model, val_dataloader, ['lung', 'heart', 'trachea']\n    )","execution_count":null,"outputs":[]},{"metadata":{"id":"J0SRSJrLaNad","outputId":"bff0ea2a-274c-48a2-920c-aa95c68502ff","trusted":true},"cell_type":"code","source":"dice_df = pd.DataFrame(dice_scores_per_classes)\ndice_df.columns = ['lung dice', 'heart dice', 'trachea dice']\n\niou_df = pd.DataFrame(iou_scores_per_classes)\niou_df.columns = ['lung jaccard', 'heart jaccard', 'trachea jaccard']\nval_metics_df = pd.concat([dice_df, iou_df], axis=1, sort=True)\nval_metics_df = val_metics_df.loc[:, ['lung dice', 'lung jaccard', \n                                      'heart dice', 'heart jaccard', \n                                      'trachea dice', 'trachea jaccard']]\nval_metics_df","execution_count":null,"outputs":[]},{"metadata":{"id":"PVbmrufYaNYh","outputId":"8785b132-f120-4421-a397-7cdc1268b07b","trusted":true},"cell_type":"code","source":"colors = ['#35FCFF', '#FF355A', '#96C503', '#C5035B', '#28B463', '#35FFAF']\npalette = sns.color_palette(colors, 6)\n\nfig, ax = plt.subplots(figsize=(12, 6));\nsns.barplot(x=val_metics_df.mean().index, y=val_metics_df.mean(), palette=palette, ax=ax);\nax.set_xticklabels(val_metics_df.columns, fontsize=14, rotation=15);\nax.set_title(\"Dice and Jaccard Coefficients from Validation\", fontsize=20)\n\nfor idx, p in enumerate(ax.patches):\n        percentage = '{:.1f}%'.format(100 * val_metics_df.mean().values[idx])\n        x = p.get_x() + p.get_width() / 2 - 0.15\n        y = p.get_y() + p.get_height()\n        ax.annotate(percentage, (x, y), fontsize=15, fontweight=\"bold\")\n\nfig.savefig(\"result1.png\", format=\"png\",  pad_inches=0.2, transparent=False, bbox_inches='tight')\nfig.savefig(\"result1.svg\", format=\"svg\",  pad_inches=0.2, transparent=False, bbox_inches='tight')","execution_count":null,"outputs":[]},{"metadata":{"id":"IxOrlMixjApG","outputId":"77687e65-d2c7-4ebd-d7b3-a1205eee4433","trusted":true},"cell_type":"code","source":"colors = ['#35FCFF', '#FF355A', '#28B463', '#35FFAF', '#96C503', '#C5035B']\npalette = sns.color_palette(colors[1::], 3)\n\nfig, ax = plt.subplots(1, 2, figsize=(20, 5))\nsns.boxplot(data=dice_df, palette=palette, ax=ax[0])\nax[0].set_ylabel(\"Dice coefficients for 3342 slices\", fontsize=14)\nax[0].set_title(\"Dice coefficients from Validation\", fontsize=20)\nax[0].set_xticklabels(dice_df.columns, fontsize=14)\n\nsns.boxplot(data=iou_df, palette=palette, ax=ax[1])\nax[1].set_ylabel(\"Jaccard coefficients for 3342 slices\", fontsize=14)\nax[1].set_title(\"Jaccard coefficients from Validation\", fontsize=20)\nax[1].set_xticklabels(iou_df.columns, fontsize=14)\nplt.tight_layout()\n\nfig.savefig(\"result2.png\", format=\"png\",  pad_inches=0.2, transparent=False, bbox_inches='tight')\nfig.savefig(\"result2.svg\", format=\"svg\",  pad_inches=0.2, transparent=False, bbox_inches='tight')","execution_count":null,"outputs":[]},{"metadata":{"id":"RCFfMqRcjaTU"},"cell_type":"markdown","source":"and train history logs "},{"metadata":{"id":"Ph3v04BCja3L","outputId":"59125857-f0c6-4af8-f855-8a1c15ec0024","trusted":true},"cell_type":"code","source":"train_logs = pd.read_csv(config.train_logs_path)\ntrain_logs.head(5)","execution_count":null,"outputs":[]},{"metadata":{"id":"aPCFHGIkja6k","outputId":"6eeaa91e-9720-49bc-ff09-35876ace4a5c","trusted":true},"cell_type":"code","source":"colors = ['#C042FF', '#03C576FF', '#FF355A', '#03C5BF', '#96C503', '#C5035B']\npalettes = [sns.color_palette(colors, 2),\n            sns.color_palette(colors, 4), \n            sns.color_palette(colors[:2]+colors[-2:] + colors[2:-2], 6)]\n            \nfig, ax = plt.subplots(1, 3, figsize=(22, 4))\n\nsns.lineplot(data=train_logs.iloc[:, :2], palette=palettes[0], markers=True, ax=ax[0], linewidth=2.5,)\nax[0].set_title(\"Loss Function during Model Training\", fontsize=14)\nax[0].set_xlabel(\"Epoch\", fontsize=14)\n\nsns.lineplot(data=train_logs.iloc[:, 2:], palette=palettes[1], markers=True, ax=ax[1], linewidth=2.5, legend=\"full\")\nax[1].set_title(\"Dice and Jaccard Coefficients during Model Training\", fontsize=14)\nax[1].set_xlabel(\"Epoch\", fontsize=14)\n\nsns.boxplot(data=val_metics_df.iloc[:,:], palette=palettes[2], ax=ax[2])\nax[2].set_title(\"Dice and Jaccard Coefficients for each Label from Validation\", fontsize=14)\nax[2].set_xticklabels(val_metics_df.columns, fontsize=10, rotation=15)\n\nplt.tight_layout()\nfig.savefig(\"result3.png\", format=\"png\",  pad_inches=0.2, transparent=False, bbox_inches='tight')\nfig.savefig(\"result3.svg\", format=\"svg\",  pad_inches=0.2, transparent=False, bbox_inches='tight')","execution_count":null,"outputs":[]},{"metadata":{"id":"v_o71Xz4j5qL"},"cell_type":"markdown","source":"Now let's make a video with overlapped masks for each slice of one id CT "},{"metadata":{"id":"jKRO64tRjAmC","outputId":"dc6b0851-5cb9-4e45-f253-232a450fc0c8","trusted":true},"cell_type":"code","source":"df = pd.read_csv(config.path_to_csv)\ndf[\"Id\"] = df['ImageId'].apply(lambda x: x.split(\"_\")[0])\n\nid_ = 'ID00400637202305055099402'\nfull_scan_example = df.loc[df['Id'] == id_].reset_index(drop=True)\nfull_scan_example ","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"get_color_info()","execution_count":null,"outputs":[]},{"metadata":{"id":"MGE1R-8ZuFw2"},"cell_type":"markdown","source":"### Ground Truth"},{"metadata":{"id":"ihOjemlbaNWA","trusted":true},"cell_type":"code","source":"PATH_TO_SAVE = id_ + \"_ground_truth\"\n\nif not os.path.exists(PATH_TO_SAVE):\n    os.mkdir(PATH_TO_SAVE)\n    print(f\"Folder {PATH_TO_SAVE} created.\")\n\nget_overlaid_masks_on_full_ctscan(ct_scan_id_df=full_scan_example,\n                                  path_to_save=\"./\" +PATH_TO_SAVE,\n                                  root_imgs_dir = '../input/chest-ct-segmentation/images/images/',        \n                                  root_masks_dir = '../input/chest-ct-segmentation/masks/masks/')\n","execution_count":null,"outputs":[]},{"metadata":{"id":"_wgi6MDjUPpJ","outputId":"7294f634-c5db-4138-8c48-f811edece77a","_kg_hide-output":false,"trusted":true},"cell_type":"code","source":"%%time\ncreate_video(path_to_imgs=PATH_TO_SAVE, video_name=id_+\"_ground_truth\", framerate=30)","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-output":true,"_kg_hide-input":false,"trusted":true},"cell_type":"code","source":"# install any fonts\n!wget -O bitwise.zip https://www.1001freefonts.com/d/8190/bitwise.zip\n!unzip bitwise.zip","execution_count":null,"outputs":[]},{"metadata":{"id":"GxDMwTMuuJ3m","_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"%%bash\n# to play in google colab had to recode\n#https://ottverse.com/ffmpeg-drawtext-filter-dynamic-overlays-timecode-scrolling-text-credits/\nffmpeg -i 'ID00400637202305055099402_ground_truth.avi' -vf \"drawtext=text='Ground Truth':x=195:y=8:fontsize=24:fontfile='./Bitwise.ttf':fontcolor='#FFFFFF'\" -strict -2 'transcoded_video1.mp4'","execution_count":null,"outputs":[]},{"metadata":{"id":"i3Sqb1kYuJ0-","outputId":"b0886607-87b2-4885-c3d3-770725b180a0","trusted":true},"cell_type":"code","source":"show_video(\"transcoded_video1.mp4\")","execution_count":null,"outputs":[]},{"metadata":{"id":"F32sZOILuXKl"},"cell_type":"markdown","source":"### Prediction"},{"metadata":{"id":"lU3dE5LmuJv6","outputId":"7908369a-a1bb-47ce-f527-cc618e1e1f71","trusted":true},"cell_type":"code","source":"imgs, predictions = get_id_predictions(net=model,\n                                       ct_scan_id_df=full_scan_example,\n                                       root_imgs_dir=config.path_to_imgs_dir)","execution_count":null,"outputs":[]},{"metadata":{"id":"SjzuZ_a9uYDJ","outputId":"62680016-1cef-4bbe-9d5c-c68e45f5e9fc","trusted":true},"cell_type":"code","source":"%%time\nPATH_TO_SAVE = id_ + \"_predictions\"\n\nif not os.path.exists(PATH_TO_SAVE):\n    os.mkdir(PATH_TO_SAVE)\n    print(f\"Folder {PATH_TO_SAVE} created.\")\n\n_= [\n    get_overlaid_masks_on_image(one_slice_image=image,\n                                one_slice_mask=mask, \n                                write=True,\n                                path_to_save=PATH_TO_SAVE,\n                                name_to_save= str(i_name)\n                                ) \n    for i_name, (image, mask) in enumerate(zip(imgs, predictions))\n    ]","execution_count":null,"outputs":[]},{"metadata":{"id":"ZFLexbXiuYIC","outputId":"5c93b9f8-089f-483f-ffd0-779249576efb","trusted":true},"cell_type":"code","source":"%%time\ncreate_video(path_to_imgs=PATH_TO_SAVE, video_name=id_+\"_predictions\", framerate=30)","execution_count":null,"outputs":[]},{"metadata":{"id":"mciRdYbIuYAo","_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"%%bash\nffmpeg -i 'ID00400637202305055099402_predictions.avi' -vf \"drawtext=text='Prediction':x=195:y=8:fontsize=24:fontfile='./Bitwise.ttf':fontcolor='#FFFFFF'\" -strict -2 'transcoded_video2.mp4'","execution_count":null,"outputs":[]},{"metadata":{"id":"IYyUAwIvulLj","outputId":"4feeb8eb-36bc-4667-f3ae-c596bca0387d","trusted":true},"cell_type":"code","source":"show_video(\"transcoded_video2.mp4\")","execution_count":null,"outputs":[]},{"metadata":{"id":"LJigaNjiuojR"},"cell_type":"markdown","source":"Merging video with ground truth slices and video with predicted slices"},{"metadata":{"id":"mY7OYIX8ulJE","_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"%%bash\n#  https://unix.stackexchange.com/questions/233832/merge-two-video-clips-into-one-placing-them-next-to-each-other\nffmpeg \\\n  -i transcoded_video1.mp4 \\\n  -i transcoded_video2.mp4 \\\n  -filter_complex '[0:v]pad=iw*2:ih[int];[int][1:v]overlay=W/2:0[vid]' \\\n  -map [vid] \\\n  -c:v libx264 \\\n  -crf 23 \\\n  -preset veryfast \\\n  result.mp4","execution_count":null,"outputs":[]},{"metadata":{"id":"ba-hcM-PuX-T","outputId":"939630c3-4f8e-45e9-a2fb-3f1ed8a98986","trusted":true},"cell_type":"code","source":"show_video(\"result.mp4\")","execution_count":null,"outputs":[]},{"metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"trusted":true},"cell_type":"code","source":"!rm -r ID00400637202305055099402_ground_truth\n!rm -r ID00400637202305055099402_predictions","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":4}