{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":65035,"databundleVersionId":7153963,"sourceType":"competition"},{"sourceId":7200393,"sourceType":"datasetVersion","datasetId":4164842}],"dockerImageVersionId":30627,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"####  This is the initial code for the competition. In this code, we handle data loading, model creation, training, and prediction. It's important to note that we have customized the segmentation models library to support a multi-task learning approach, as detailed in our research paper","metadata":{}},{"cell_type":"markdown","source":"#### Install required libraries to execute this notebook","metadata":{}},{"cell_type":"code","source":"!pip install split-folders[full] natsort timm==0.6.12 pretrainedmodels==0.7.4 efficientnet_pytorch==0.7.1 albumentations==1.3.1","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:26.706587Z","iopub.execute_input":"2023-12-14T13:49:26.707143Z","iopub.status.idle":"2023-12-14T13:49:39.190503Z","shell.execute_reply.started":"2023-12-14T13:49:26.707100Z","shell.execute_reply":"2023-12-14T13:49:39.189307Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"","metadata":{}},{"cell_type":"code","source":"import shutil\n# copying the train data to the kaggle working direcrtory\nshutil.copytree('/kaggle/input/breast-density-prediction/train', '/kaggle/working', dirs_exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:39.193446Z","iopub.execute_input":"2023-12-14T13:49:39.193900Z","iopub.status.idle":"2023-12-14T13:49:43.396812Z","shell.execute_reply.started":"2023-12-14T13:49:39.193856Z","shell.execute_reply":"2023-12-14T13:49:43.395759Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\n\n# Source path (your uploaded library folder)\nsource_path = '/kaggle/input/mtlsegnet'\n# Destination path (Kaggle notebook's working directory)\ndestination_path = '/kaggle/working'\n# Copy the contents of the source folder to the destination\nshutil.copytree(source_path, destination_path, dirs_exist_ok=True)\n","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:43.398046Z","iopub.execute_input":"2023-12-14T13:49:43.398400Z","iopub.status.idle":"2023-12-14T13:49:43.554287Z","shell.execute_reply.started":"2023-12-14T13:49:43.398370Z","shell.execute_reply":"2023-12-14T13:49:43.553305Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs('/kaggle/working/data', exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:43.555573Z","iopub.execute_input":"2023-12-14T13:49:43.555885Z","iopub.status.idle":"2023-12-14T13:49:43.560661Z","shell.execute_reply.started":"2023-12-14T13:49:43.555859Z","shell.execute_reply":"2023-12-14T13:49:43.559588Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Splitting the data for training (80%) and validation (20%)","metadata":{}},{"cell_type":"code","source":"import splitfolders\n\n# Split with a ratio.\n# To only split into training and validation set, set a tuple to `ratio`, i.e, `(.8, .2)`.\nsplitfolders.ratio(\"/kaggle/working/train\", \n                   output=\"/kaggle/working/data\",\n                   seed=1337, ratio=(.8, .2), group_prefix=None, move=False) # default values","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:43.563243Z","iopub.execute_input":"2023-12-14T13:49:43.563578Z","iopub.status.idle":"2023-12-14T13:49:44.404479Z","shell.execute_reply.started":"2023-12-14T13:49:43.563554Z","shell.execute_reply":"2023-12-14T13:49:44.403547Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Import required libraries\n\nimport os\nimport glob\nfrom natsort import natsorted\nimport numpy as np\nfrom PIL import Image\nimport cv2\nimport torch\nfrom torch.utils.data import Dataset\nfrom torchvision import transforms\nimport albumentations as A\nimport torch\nimport torch.optim as optim\nimport torch.nn as nn\nfrom torch.utils.data import DataLoader\nimport segmentation_models_multi_tasking as smp\nimport matplotlib.pyplot as plt\n\n","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:44.405827Z","iopub.execute_input":"2023-12-14T13:49:44.406113Z","iopub.status.idle":"2023-12-14T13:49:44.411941Z","shell.execute_reply.started":"2023-12-14T13:49:44.406088Z","shell.execute_reply":"2023-12-14T13:49:44.411079Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"##### Creating data loader function ","metadata":{}},{"cell_type":"code","source":"class MammoDataset(Dataset):\n    def __init__(self, path, dataset, split, augmentations=False):\n        self.path = path\n        self.split = split\n        self.dataset = dataset\n        self.augmentations = augmentations\n\n        self.images = natsorted(glob.glob(os.path.join(self.path, self.dataset, self.split, 'images/*')))\n        self.masks = natsorted(glob.glob(os.path.join(self.path, self.dataset, self.split, 'breast_masks/*')))\n        self.contours = natsorted(glob.glob(os.path.join(self.path, self.dataset, self.split, 'dense_masks/*')))\n        \n    def __len__(self):\n        return len(self.images)\n    \n    def __getitem__(self, index):\n\n        self.to_tensor = transforms.Compose([transforms.ToPILImage(),\n                                             transforms.Resize((256, 256)),\n                                             transforms.ToTensor(),])\n\n        self.aug_pipeline = A.Compose([A.ShiftScaleRotate(p = 0.5),\n                                       A.HorizontalFlip(0.5),\n                                       A.VerticalFlip(0.5),\n                                       A.RandomBrightnessContrast(0.5),\n                                       A.GaussNoise(p = 0.5),\n                                       A.GaussianBlur(p=0.5),\n                                       A.ElasticTransform(),],p=0.8)\n\n\n        self.image = cv2.imread(self.images[index], 1)\n        self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)\n\n        self.mask = cv2.imread(self.masks[index], 0)\n        self.contour = cv2.imread(self.contours[index], 0)\n\n        if self.augmentations:\n            torch.manual_seed(1990)\n            masks = [self.mask, self.contour]\n            #if self.augmentation_type == 'pixel':\n            #    self.sample = self.pixel_augmentations(image = self.image, masks = masks)\n            #elif self.augmentation_type == 'spatial':\n            #    self.sample = self.spatial_augmentations(image=self.image, masks=masks)\n            #elif self.augmentation_type == 'pixel_spatial':\n            #    self.sample = self.pixel_spatial_augmentations(image=self.image, masks=masks)\n            #elif self.augmentation_type == 'aug_pipeline':\n            self.sample = self.aug_pipeline(image=self.image, masks=masks)\n\n            return self.to_tensor(self.sample['image']), self.to_tensor(self.sample['masks'][0]), self.to_tensor(self.sample['masks'][1])\n        else:\n            return self.to_tensor(self.image), self.to_tensor(self.mask), self.to_tensor(self.contour)\n","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:44.413097Z","iopub.execute_input":"2023-12-14T13:49:44.413363Z","iopub.status.idle":"2023-12-14T13:49:47.824379Z","shell.execute_reply.started":"2023-12-14T13:49:44.413339Z","shell.execute_reply":"2023-12-14T13:49:47.823228Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Loading the train and valid dataset","metadata":{}},{"cell_type":"code","source":"# loading training dataset\ntrain_dataset = MammoDataset(path='/kaggle/working', dataset='data', split='train', augmentations=None)\nprint(len(train_dataset))\ntrain_dataloader = DataLoader(train_dataset, shuffle=True, batch_size =4, num_workers=0)\n\n# loading valid dataset\nvalid_dataset = MammoDataset(path='/kaggle/working', dataset='data', split='val', augmentations=None)\nprint(len(valid_dataset))\nvalid_dataloader = DataLoader(valid_dataset, shuffle=True, batch_size=4, num_workers=0)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:47.825661Z","iopub.execute_input":"2023-12-14T13:49:47.825964Z","iopub.status.idle":"2023-12-14T13:49:47.914519Z","shell.execute_reply.started":"2023-12-14T13:49:47.825938Z","shell.execute_reply":"2023-12-14T13:49:47.913575Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# device to work (CPU/GPU)\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nDEVICE","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:47.915700Z","iopub.execute_input":"2023-12-14T13:49:47.916052Z","iopub.status.idle":"2023-12-14T13:49:47.922462Z","shell.execute_reply.started":"2023-12-14T13:49:47.916025Z","shell.execute_reply":"2023-12-14T13:49:47.921582Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# create segmentation model with pretrained encoder, here we are using Unet architecrure with Resnet50 as encoder with ImageNet pretrained weights\nmodel = getattr(smp, 'Unet')(\n    encoder_name= 'resnet50',\n    encoder_weights = 'imagenet', \n    classes =1,\n    activation = 'sigmoid')\n\nmodel = model.to(DEVICE)\nmodel = nn.DataParallel(model)\nprint(model, (3, 256, 256))","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:47.923612Z","iopub.execute_input":"2023-12-14T13:49:47.923967Z","iopub.status.idle":"2023-12-14T13:49:49.948516Z","shell.execute_reply.started":"2023-12-14T13:49:47.923935Z","shell.execute_reply":"2023-12-14T13:49:49.947571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# We modified the FocalTverskyLoss function to adopt the multi-task learning approach proposed in our paper\nloss = getattr(smp.utils.losses, 'FocalTverskyLoss')()\n\n# Adam optimizer with reduce learning rate schedular\noptimizer = getattr(torch.optim, 'Adam')([dict(params=model.parameters(), lr=0.0001),])\nlr_schedular = optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.25, patience=2)\n\n# Evaluation metrics\nmetrics = [\n    smp.utils.metrics.Precision(),\n    smp.utils.metrics.Recall(),\n    smp.utils.metrics.Accuracy(),\n    smp.utils.metrics.Fscore(),\n    smp.utils.metrics.IoU(threshold=0.5),\n    \n]","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:49.949755Z","iopub.execute_input":"2023-12-14T13:49:49.950088Z","iopub.status.idle":"2023-12-14T13:49:49.959025Z","shell.execute_reply.started":"2023-12-14T13:49:49.950059Z","shell.execute_reply":"2023-12-14T13:49:49.958064Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# 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    lr_schedular=lr_schedular,\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)","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:49.960191Z","iopub.execute_input":"2023-12-14T13:49:49.960465Z","iopub.status.idle":"2023-12-14T13:49:49.983922Z","shell.execute_reply.started":"2023-12-14T13:49:49.960440Z","shell.execute_reply":"2023-12-14T13:49:49.983035Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.makedirs(\"/kaggle/working/results/logs\", exist_ok=True) # saving the logs \nos.makedirs(\"/kaggle/working/results/models\", exist_ok=True) # asving the trained model\n\ntrain_accuracy = []\nvalid_accuracy = []\ntrain_loss = []\nvalid_loss = []\n\nnum_epochs = 30\n\nwith open(\"/kaggle/working/results/logs/unet_resnet50_logs.txt\", 'a+') as logs_file:\n    # train model for 40 epochs\n    \n    max_score = 0\n    for i in range(0, num_epochs):# 30 epochs\n        print('\\nEpoch: {}'.format(i))\n        train_logs = train_epoch.run(train_dataloader)\n        valid_logs = valid_epoch.run(valid_dataloader)\n\n\n        train_accuracy.append(train_logs['accuracy'])\n        valid_accuracy.append(valid_logs['accuracy'])\n        train_loss.append(train_logs['focal_tversky_loss_weighted'])\n        valid_loss.append(valid_logs['focal_tversky_loss_weighted'])\n\n        print('{} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {} \\t {}'.format(i, train_logs['focal_tversky_loss_breast'],\n                                                                                                                    train_logs['focal_tversky_loss_dense'],\n                                                                                                                    train_logs['focal_tversky_loss_weighted'],\n                                                                                                                    train_logs['precision'],\n                                                                                                                    train_logs['recall'],\n                                                                                                                    train_logs['accuracy'],\n                                                                                                                    train_logs['fscore'],\n                                                                                                                    train_logs['iou_score'],\n                                                                                                                    valid_logs['focal_tversky_loss_breast'],\n                                                                                                                    valid_logs['focal_tversky_loss_dense'],\n                                                                                                                    valid_logs['focal_tversky_loss_weighted'],\n                                                                                                                    valid_logs['precision'],\n                                                                                                                    valid_logs['recall'],\n                                                                                                                    valid_logs['accuracy'],\n                                                                                                                    valid_logs['fscore'],\n                                                                                                                    valid_logs['iou_score']), file=logs_file)\n\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, \"/kaggle/working/results/models/unet_resnet50_model.pth\")\n            print('Model saved!')","metadata":{"execution":{"iopub.status.busy":"2023-12-14T13:49:49.985251Z","iopub.execute_input":"2023-12-14T13:49:49.985493Z","iopub.status.idle":"2023-12-14T15:10:03.860960Z","shell.execute_reply.started":"2023-12-14T13:49:49.985472Z","shell.execute_reply":"2023-12-14T15:10:03.860033Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"##### Utility codes for plotting the learning curves, to visualize the segmentation results, and to predict the density","metadata":{}},{"cell_type":"code","source":"# Convert the image to a pytorch tensor\ndef image_tensor(img):\n    \"\"\"\n    Converts a PIL Image or NumPy array to a PyTorch tensor.\n    Args:\n    - img: PIL.Image or np.ndarray, the image to be converted to a PyTorch tensor.\n    Returns:\n    - image: torch.Tensor, the converted PyTorch tensor.\n    Raises:\n    - TypeError: if the input image is not a PIL.Image or np.ndarray.\n    \"\"\"\n    \n    if type(img) not in [np.ndarray, Image.Image]:\n        raise TypeError(\"Input must be np.ndarray or PIL.Image\")\n\n    # Define a PyTorch tensor transformer pipeline\n    torch_tensor = transforms.Compose(\n        [transforms.Resize((256, 256)), transforms.ToTensor()]\n    )\n\n    if type(img) == Image.Image:\n        # Convert PIL image to PyTorch tensor\n        image = torch_tensor(img)\n        image = image.unsqueeze(0)\n        print(\"tensor\", type(image))\n        return image\n    elif type(img) == np.ndarray:\n        # Convert NumPy array to RGB PIL image and then to PyTorch tensor\n        pil_image = Image.fromarray(img).convert(\"RGB\")\n        image = torch_tensor(pil_image)\n        image = image.unsqueeze(0)\n        print(\"tensor\", type(image))\n        return image\n    else:\n        raise TypeError(\"Input must be np.ndarray or PIL.Image\")\n\n\ndef mask_to_rgba(mask, color=\"red\"):\n    MASK_COLORS = [\"red\", \"green\", \"blue\", \"yellow\", \"magenta\", \"cyan\"]\n    assert color in MASK_COLORS\n    assert mask.ndim == 3 or mask.ndim == 2\n\n    h = mask.shape[0]\n    w = mask.shape[1]\n    zeros = np.zeros((h, w))\n    ones = mask.reshape(h, w)\n    if color == \"red\":\n        return np.stack((ones, zeros, zeros, ones), axis=-1)\n    elif color == \"green\":\n        return np.stack((zeros, ones, zeros, ones), axis=-1)\n    elif color == \"blue\":\n        return np.stack((zeros, zeros, ones, ones), axis=-1)\n    elif color == \"yellow\":\n        return np.stack((ones, ones, zeros, ones), axis=-1)\n    elif color == \"magenta\":\n        return np.stack((ones, zeros, ones, ones), axis=-1)\n    elif color == \"cyan\":\n        return np.stack((zeros, ones, ones, ones), axis=-1)\n\n    \n# learning curves visualization        \ndef plot_learning_curves(train_accuracy, valid_accuracy, train_loss, valid_loss, num_epochs):\n    \"\"\"\n    Plots the learning curves for training and validation accuracy and loss.\n    \n    Parameters:\n    - train_accuracy: List of training accuracy values for each epoch.\n    - valid_accuracy: List of validation accuracy values for each epoch.\n    - train_loss: List of training loss values for each epoch.\n    - valid_loss: List of validation loss values for each epoch.\n    - epochs: List of epoch numbers.\n    \"\"\"\n    fig, axes = plt.subplots(1, 2, figsize=(12, 4))  # 1 row, 2 columns for subplots\n    \n    epochs = range(1, num_epochs+1)\n    # Plot accuracy on the first subplot\n    axes[0].plot(epochs, train_accuracy, label='Train Accuracy')\n    axes[0].plot(epochs, valid_accuracy, label='Validation Accuracy')\n    axes[0].set_xlabel('Epochs')\n    axes[0].set_ylabel('Accuracy')\n    axes[0].set_title('Accuracy Curve')\n    axes[0].legend()\n\n    # Plot loss on the second subplot\n    axes[1].plot(epochs, train_loss, label='Train Loss')\n    axes[1].plot(epochs, valid_loss, label='Validation Loss')\n    axes[1].set_xlabel('Epochs')\n    axes[1].set_ylabel('Loss')\n    axes[1].set_title('Loss Curve')\n    axes[1].legend()\n\n    # Adjust spacing between subplots\n    plt.tight_layout()\n\n    # Display the subplots\n    plt.show()\n\n# Visualize the breast and dense tissue segmentation\ndef visualize_segmentation_results(image_path, model_path):\n    # Load and preprocess the image\n    image = Image.open(image_path).convert('RGB')\n    img = image_tensor(image)\n\n    # Load the model\n    model = torch.load(model_path)\n    model.eval()\n\n    # Make predictions using the model\n    with torch.no_grad():\n        pred1, pred2 = model.module.predict(img.cuda())\n\n    # Convert predictions to numpy arrays\n    img = img[0].cpu().numpy().transpose(1, 2, 0)\n    img = img[:, :, 0]\n\n    pred1 = pred1[0].cpu().numpy().transpose(1, 2, 0)\n    pred1 = pred1[:, :, 0]\n\n    pred2 = pred2[0].cpu().numpy().transpose(1, 2, 0)\n    pred2 = pred2[:, :, 0]\n\n    fig, axes = plt.subplots(1, 2, figsize=(20, 15), squeeze=False)\n\n    axes[0, 0].set_title('Original image', fontsize=25)\n    axes[0, 1].set_title(\"Breast and dense tissue segmentation\", fontsize=25)\n\n    axes[0, 0].imshow(img, cmap='gray')\n    axes[0, 0].set_axis_off()\n\n    axes[0, 1].imshow(img, cmap='gray')\n    axes[0, 1].imshow(mask_to_rgba(pred1, color='red'), cmap='gray', alpha=0.2)\n    axes[0, 1].imshow(mask_to_rgba(pred2, color='green'), cmap='gray', alpha=0.5)\n    axes[0, 1].set_axis_off()\n","metadata":{"execution":{"iopub.status.busy":"2023-12-14T15:10:03.863955Z","iopub.execute_input":"2023-12-14T15:10:03.864248Z","iopub.status.idle":"2023-12-14T15:10:03.887907Z","shell.execute_reply.started":"2023-12-14T15:10:03.864223Z","shell.execute_reply":"2023-12-14T15:10:03.886915Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Learning curves visualization\nplot_learning_curves(train_accuracy, valid_accuracy, train_loss, valid_loss, num_epochs)\n","metadata":{"execution":{"iopub.status.busy":"2023-12-14T15:10:03.889245Z","iopub.execute_input":"2023-12-14T15:10:03.889637Z","iopub.status.idle":"2023-12-14T15:10:04.549816Z","shell.execute_reply.started":"2023-12-14T15:10:03.889603Z","shell.execute_reply":"2023-12-14T15:10:04.548822Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"visualize_segmentation_results(image_path='/kaggle/working/data/train/images/034d4806fba1ea27cb26a890364a30ce.jpg', \n                               model_path=\"/kaggle/working/results/models/unet_resnet50_model.pth\")","metadata":{"execution":{"iopub.status.busy":"2023-12-14T15:10:04.550990Z","iopub.execute_input":"2023-12-14T15:10:04.551275Z","iopub.status.idle":"2023-12-14T15:10:05.909697Z","shell.execute_reply.started":"2023-12-14T15:10:04.551249Z","shell.execute_reply":"2023-12-14T15:10:05.908685Z"},"trusted":true},"execution_count":null,"outputs":[]}]}