0.0-rosia-segmenting-subsurface.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🏆 Every Layer, Everywhere, All at Once: Segmenting Subsurface\n",
    "\n",
    "<img src='assets/banner-logo.png'>\n",
    "\n",
    "## 🌋 Context\n",
    "\n",
    "The competition challenges participants to utilize Meta's Segment Anything Model (SAM) for new use cases beyond traditional segmentation tasks in 3D seismic data analysis. The primary objective is to develop solutions capable of identifying and mapping all layers within seismic data simultaneously. By doing so, participants aim to accelerate the interpretation process, enabling quicker analysis of large datasets and fostering a deeper understanding of Earth's structure and geological features. The provided dataset consists of approximately 9,000 pre-interpreted seismic volumes, each accompanied by segment masks for model training. These volumes represent diverse geological settings and present typical challenges of seismic data interpretation, including complex geology and data processing workflows. The holdout data for evaluation mirrors the complexity of the training data, ensuring robust solutions capable of handling diverse geologic features across different seismic volumes.\n",
    "\n",
    "## 🧠 Our approach\n",
    "\n",
    "<img src='assets/approach.png'>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📚 Libraries\n",
    "\n",
    "Our code run on `Python 3.10.13`\n",
    "\n",
    "Installing external libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Installing segmenting-subsurface librarie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -e ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generic\n",
    "import os\n",
    "\n",
    "# Data manipulation\n",
    "import numpy as np\n",
    "\n",
    "# Visualization\n",
    "import torch\n",
    "import torchvision.transforms.functional as tvF\n",
    "import plotly.graph_objects as go\n",
    "from plotly.subplots import make_subplots\n",
    "\n",
    "# Custom packages (accessible into src directory)\n",
    "import src.data.make_dataset as md\n",
    "import src.features.mask2former_inference as mask2former_inf\n",
    "import src.features.segformer_inference as segformer_inf\n",
    "import src.features.segment_anything_inference as sam_inf\n",
    "import src.visualization.utils as vutils\n",
    "from src import utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Connecting to RosIA wandb demonstration account"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# W&B initialisation\n",
    "wandb_api_key = '02ca932e1203e93aaa8c97b8331d6c0b04c3170a'  # Please do not share this api key\n",
    "! wandb login --relogin {wandb_api_key}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 📸 Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "volume_path = 'data/raw/train/69764103/seismic_block-2_vol_69764103.npy'\n",
    "volume = np.load(volume_path, allow_pickle=True)\n",
    "\n",
    "label_path = volume_path.replace('seismic', 'horizon_labels')\n",
    "label = np.load(label_path, allow_pickle=True)\n",
    "\n",
    "volume_hollow = vutils.get_volume_hollow(volume)\n",
    "volume_plotly = vutils.get_plotly_volume(volume_hollow, colorscale='Greys')\n",
    "\n",
    "label_hollow = vutils.get_volume_hollow(label)\n",
    "label_plotly = vutils.get_plotly_volume(label_hollow, colorscale='Viridis')\n",
    "\n",
    "fig = make_subplots(\n",
    "    rows=1, cols=2,\n",
    "    specs=[\n",
    "        [{'type': 'scatter3d'}, {'type': 'scatter3d'}]\n",
    "    ],\n",
    "    subplot_titles=('Original volume', 'Labelised volume')\n",
    ")\n",
    "\n",
    "fig.add_trace(volume_plotly, row=1, col=1)\n",
    "fig.add_trace(label_plotly, row=1, col=2)\n",
    "fig.update_layout(showlegend=False)\n",
    "fig.show()\n",
    "\n",
    "del label_hollow, volume_hollow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ⚒️ Preprocessing\n",
    "\n",
    "In computer vision, the classic preprocessing steps for an image are as follows:\n",
    "\n",
    "1. `Scaling`: Allows us to scale the values between 0 and 1. (Using a Min Max Scaler)\n",
    "\n",
    "2. `Normalization`: Helps us achieve a Gaussian distribution of values for each channel. (Using a Standard Scaler)\n",
    "\n",
    "3. `Rescaling`: If necessary, based on what input the model accepts. (Using a bilinear interpolation)\n",
    "\n",
    "However, we have observed that in the images of our dataset, the objective is to delineate areas of varying brightness between them. That's why we decided to add `contrast` to highlight these differences in shade between the layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = utils.get_config()\n",
    "\n",
    "\n",
    "def scale(image):\n",
    "    image = (image - config['data']['min']) / (config['data']['max'] - config['data']['min'])\n",
    "    return image\n",
    "\n",
    "\n",
    "def contrast(image):\n",
    "    tensor = torch.from_numpy(image).unsqueeze(0)\n",
    "    tensor = tvF.adjust_contrast(tensor, contrast_factor=25)\n",
    "\n",
    "    return tensor.squeeze().numpy(force=True)\n",
    "\n",
    "\n",
    "def normalize(image):\n",
    "    tensor = torch.from_numpy(image).unsqueeze(0)\n",
    "    tensor = tvF.normalize(tensor, mean=[config['data']['mean']], std=[config['data']['std']])\n",
    "\n",
    "    return tensor.squeeze().numpy(force=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "slice_idx = 200\n",
    "\n",
    "image = volume[slice_idx, :, :].T\n",
    "\n",
    "fig = make_subplots(\n",
    "    rows=2, cols=2,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}, {'type': 'heatmap'}],\n",
    "        [{'type': 'heatmap'}, {'type': 'heatmap'}]\n",
    "    ],\n",
    "    subplot_titles=('0 - Original image', '1 - Scaled image', '2 - Contrasted image', '3 - Normalized image')\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys'), row=1, col=1)\n",
    "image = scale(image)\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys'), row=1, col=2)\n",
    "image = contrast(image)\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys'), row=2, col=1)\n",
    "image = normalize(image)\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys'), row=2, col=2)\n",
    "fig.update_layout(showlegend=False)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🤖 Segment Anything Model (SAM): Binary mask\n",
    "\n",
    "We initially attempted to predict all the classes for each image by assuming that each label corresponded to a distinct rock layer. In other words, we placed ourselves in an instance segmentation problem with as many classes as distinct labels, and we fine-tuned a Mask2former model, but this approach did not yield good results.\n",
    "\n",
    "Therefore, we pursued another approach. Since instance segmentation doesn't seem to produce results, we simplified the problem by converting the labels to 0 or 1 based on their parity. We transformed our problem into binary semantic segmentation. For this, we used a model from the [Segment Anything Model](https://arxiv.org/pdf/2304.02643.pdf) family."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_binary_label(label):\n",
    "    binary_label = np.where(label % 2 == 0, 1, 0)\n",
    "\n",
    "    return binary_label\n",
    "\n",
    "fig = make_subplots(\n",
    "    rows=1, cols=3,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}] * 3\n",
    "    ],\n",
    "    subplot_titles=('Processed image', 'Original label', 'Binarized label')\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=1)\n",
    "image_label = label[slice_idx, :, :].T\n",
    "fig.add_trace(go.Heatmap(z=image_label.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=2)\n",
    "binary_label = get_binary_label(image_label.copy())\n",
    "fig.add_trace(go.Heatmap(z=binary_label.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=3)\n",
    "fig.update_layout(showlegend=False)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the SAM inference class to make predictions on the volume.\n",
    "\n",
    "sam_inference = sam_inf.SAMInference(\n",
    "    config=config,\n",
    "    cuda_idx=0,\n",
    "    list_volume=[volume_path],\n",
    "    run=None,\n",
    "    split='train',\n",
    "    batch=5  # Reduce the batch size if you encounter cuda memory issues running the inference. Min 2 - Max 300\n",
    ")\n",
    "\n",
    "volume_name = os.path.basename(volume_path)\n",
    "binary_mask_path = sam_inference.get_mask_path(volume_name)\n",
    "sam_dir = os.path.split(binary_mask_path)[0]\n",
    "os.makedirs(sam_dir, exist_ok=True)\n",
    "sam_inference()\n",
    "binary_mask = np.load(binary_mask_path, allow_pickle=True)\n",
    "sam_binary_pred = binary_mask[slice_idx, :, :].T.astype(np.uint8)\n",
    "\n",
    "# Deleting the instance to free up memory space.\n",
    "del sam_inference, binary_mask\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = make_subplots(\n",
    "    rows=1, cols=3,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}] * 3\n",
    "    ],\n",
    "    subplot_titles=('Processed image', 'Binarized label', 'SAM Prediction')\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=1)\n",
    "fig.add_trace(go.Heatmap(z=binary_label.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=2)\n",
    "fig.add_trace(go.Heatmap(z=sam_binary_pred.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=3)\n",
    "fig.update_layout(showlegend=False)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "💡 We used SAM without fine-tuning because it provides sufficient performances to extract patterns from the slices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🏞️ Segformer: Binary mask\n",
    "\n",
    "We refined the SAM predictions by using another model to generate binary masks. For this, we used a model from the [Segformer](https://proceedings.neurips.cc/paper/2021/file/64f1f27bf1b4ec22924fd0acb550c235-Paper.pdf) family."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = make_subplots(\n",
    "    rows=1, cols=3,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}] * 3\n",
    "    ],\n",
    "    subplot_titles=(\n",
    "    'Red channel (Processed image)', 'Green channel (SAM prediction)', 'Blue channel (Processed image)'),\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=1)\n",
    "fig.add_trace(go.Heatmap(z=sam_binary_pred.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=2)\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=3)\n",
    "fig.update_layout(showlegend=False, title='Decomposition of the Segformer input image into channels (RGB)')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the Segformer inference class to make predictions on the volume.\n",
    "\n",
    "run = utils.RunDemo('segformer.yml', id='mmw4795a', name='abundant-lantern-1231')\n",
    "\n",
    "segformer_inference = segformer_inf.SegformerInference(\n",
    "    config=config,\n",
    "    cuda_idx=0,\n",
    "    list_volume=[volume_path],\n",
    "    run=run,\n",
    "    split='train',\n",
    "    batch=50  # Reduce the batch size if you encounter issues running the inference. Min 2 - Max 300\n",
    ")\n",
    "\n",
    "volume_name = os.path.basename(volume_path)\n",
    "binary_mask_path = segformer_inference.get_mask_path(volume_name)\n",
    "segformer_dir = os.path.split(binary_mask_path)[0]\n",
    "os.makedirs(segformer_dir, exist_ok=True)\n",
    "segformer_inference()\n",
    "binary_mask = np.load(binary_mask_path, allow_pickle=True)\n",
    "seg_binary_pred = binary_mask[slice_idx, :, :].T.astype(np.uint8)\n",
    "\n",
    "# Deleting the instance to free up memory space.\n",
    "del segformer_inference, binary_mask\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = make_subplots(\n",
    "    rows=1, cols=3,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}] * 3\n",
    "    ],\n",
    "    subplot_titles=('Processed image', 'Binarized label', 'Segformer prediction')\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=1)\n",
    "fig.add_trace(go.Heatmap(z=binary_label.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=2)\n",
    "fig.add_trace(go.Heatmap(z=seg_binary_pred.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=3)\n",
    "fig.update_layout(showlegend=False)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With this approach, we achieved a DICE score of `0.8408` and a Jaccard index (IoU) of `0.7523` on the binary masks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🎭 Mask2Former: Instance mask\n",
    "\n",
    "Once we have obtained our binary mask, we need to obtain each layer independently of its parity in order to create a prompt as precise as possible for the Segment Anything model.\n",
    "\n",
    "For this purpose, we will use another segmentation model called [Mask2Former](http://openaccess.thecvf.com/content/CVPR2022/papers/Cheng_Masked-Attention_Mask_Transformer_for_Universal_Image_Segmentation_CVPR_2022_paper.pdf). This model will take as input the combination of the binary mask from Segformer and the original image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_instance_label(label):\n",
    "    instance_label = np.full(label.shape, np.nan)\n",
    "    old_labels = np.unique(label)\n",
    "    new_labels = range(len(old_labels))\n",
    "    for old_label, new_label in zip(old_labels, new_labels):\n",
    "        instance_label = np.where(label == old_label, new_label, instance_label)\n",
    "\n",
    "    return instance_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = make_subplots(\n",
    "    rows=1, cols=3,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}] * 3\n",
    "    ],\n",
    "    subplot_titles=('Red channel (Seg prediction)', 'Green channel (Original image)', 'Blue channel (Seg prediction)'),\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=seg_binary_pred.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=1)\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=2)\n",
    "fig.add_trace(go.Heatmap(z=seg_binary_pred.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=3)\n",
    "fig.update_layout(showlegend=False, title='Decomposition of the Mask2Former input image into channels (RGB)')\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Its objective will be to predict a variant of the original label. The IDs of the original masks are set to 0 up to the number of masks present in the original label."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the Mask2Former inference class to make predictions on the volume.\n",
    "\n",
    "run = utils.RunDemo('mask2former.yml', id='nvbtr9k2', name='vermilion-moon-1241')\n",
    "\n",
    "mask2former_inference = mask2former_inf.Mask2formerInference(\n",
    "    config=config,\n",
    "    cuda_idx=0,\n",
    "    list_volume=[volume_path],\n",
    "    run=run,\n",
    "    split='train',\n",
    "    batch=25  # Reduce the batch size if you encounter cuda memory issues running the inference. Min 2 - Max 300\n",
    ")\n",
    "\n",
    "volume_name = os.path.basename(volume_path)\n",
    "instance_mask_path = mask2former_inference.get_mask_path(volume_name)\n",
    "mask2former_dir = os.path.split(instance_mask_path)[0]\n",
    "os.makedirs(mask2former_dir, exist_ok=True)\n",
    "\n",
    "mask2former_inference()\n",
    "instance_mask = np.load(instance_mask_path, allow_pickle=True)\n",
    "intance_pred = instance_mask[slice_idx, :, :].T\n",
    "\n",
    "# Deleting the instance to free up memory space.\n",
    "del mask2former_inference, instance_mask\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image_label = label[slice_idx, :, :].T\n",
    "variant_label = get_instance_label(image_label)\n",
    "\n",
    "fig = make_subplots(\n",
    "    rows=1, cols=3,\n",
    "    specs=[\n",
    "        [{'type': 'heatmap'}] * 3\n",
    "    ],\n",
    "    subplot_titles=('Original image', 'Variant label', 'Mask2Former prediction')\n",
    ")\n",
    "\n",
    "fig.add_trace(go.Heatmap(z=image.tolist(), showscale=False, colorscale='Greys', hoverinfo='skip'), row=1, col=1)\n",
    "fig.add_trace(go.Heatmap(z=variant_label.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=2)\n",
    "fig.add_trace(go.Heatmap(z=intance_pred.tolist(), showscale=False, colorscale='viridis', hoverinfo='skip'), row=1, col=3)\n",
    "fig.update_layout(showlegend=False)\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we do not want any class or background, we have chosen to set the weights of the focus loss and the cross-entropy loss to 1, while the dice loss has a weight of 10. Refer to [Hugging Face](https://huggingface.co/docs/transformers/v4.37.2/en/model_doc/mask2former#transformers.Mask2FormerForUniversalSegmentation).\n",
    "\n",
    "With this approach, we achieved a validation loss of `22.575` on the instance masks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 🧑🏻‍💻 Code Submission"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to change the configuration of the models, please refer to the YAML file available in the config folder.\n",
    "\n",
    "The script takes volumes from the data/raw/train folder for training and data/raw/test for inference. A directory data/processed/[train, test]/{run_id} is created for all intermediate masks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# W&B initialisation\n",
    "wandb_api_key = '02ca932e1203e93aaa8c97b8331d6c0b04c3170a'  # Please do not share this api key\n",
    "! wandb login --relogin {wandb_api_key}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src import utils\n",
    "import src.models.segformer.train_model as segformer_tm\n",
    "import src.models.mask2former.train_model as mask2former_tm\n",
    "import src.features.segment_anything_inference as sam_inf\n",
    "import src.features.segformer_inference as segformer_inf\n",
    "import src.features.mask2former_inference as mask2former_inf\n",
    "import src.data.make_dataset as md\n",
    "import wandb\n",
    "import os\n",
    "import torch\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings('ignore')\n",
    "torch.set_float32_matmul_precision('medium')\n",
    "\n",
    "\n",
    "class SegmentationPipeline:\n",
    "    def __init__(self, sam_batch, segformer_batch, mask2former_batch) -> None:\n",
    "        self.config = utils.get_config()\n",
    "        self.sam_batch = sam_batch\n",
    "        self.segformer_batch = segformer_batch\n",
    "        self.mask2former_batch = mask2former_batch\n",
    "\n",
    "    def make_sam_inference(self, split):\n",
    "        sam_inference = sam_inf.SAMInference(\n",
    "            config=self.config,\n",
    "            cuda_idx=0,\n",
    "            list_volume=md.get_volumes(self.config, split),\n",
    "            run=None,\n",
    "            split=split,\n",
    "            batch=self.sam_batch\n",
    "        )\n",
    "        \n",
    "        os.makedirs(sam_inference.get_folder_path(), exist_ok=True)\n",
    "        sam_inference()\n",
    "        del sam_inference\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "    def make_segformer_inference(self, split, segformer_id):\n",
    "        segformer_inference = segformer_inf.SegformerInference(\n",
    "            config=self.config,\n",
    "            cuda_idx=0,\n",
    "            list_volume=md.get_volumes(self.config, split),\n",
    "            run=utils.RunDemo('segformer.yml', **segformer_id),\n",
    "            split=split,\n",
    "            batch=self.sam_batch\n",
    "        )\n",
    "        \n",
    "        os.makedirs(segformer_inference.get_folder_path(), exist_ok=True)\n",
    "        segformer_inference()\n",
    "        del segformer_inference\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "    def make_mask2former_inference(self, split, mask2former_id, segformer_id):\n",
    "        run = utils.RunDemo('mask2former.yml', **mask2former_id)\n",
    "        run.config['model_mask_id'] = f'{segformer_id[\"name\"]}-{segformer_id[\"id\"]}'\n",
    "        mask2former_inference = mask2former_inf.Mask2formerInference(\n",
    "            config=self.config,\n",
    "            cuda_idx=0,\n",
    "            list_volume=md.get_volumes(self.config, split),\n",
    "            run=run,\n",
    "            split=split,\n",
    "            batch=self.sam_batch\n",
    "        )\n",
    "        \n",
    "        os.makedirs(mask2former_inference.get_folder_path(), exist_ok=True)\n",
    "        mask2former_inference()\n",
    "        \n",
    "        # Deleting the instance to free up memory space.\n",
    "        del mask2former_inference\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "    def train(self):\n",
    "        print('Create Segment Anything binary masks...')\n",
    "        self.make_sam_inference(split='train')\n",
    "        print('Segment Anything binary masks done!')\n",
    "        print('Train Segformer...')\n",
    "        segformer_id = self.train_segformer()\n",
    "        print('Segformer training finished!')\n",
    "        print('Create Segformer binary masks...')\n",
    "        self.make_segformer_inference(split='train', segformer_id=segformer_id)\n",
    "        print('Segformer binary masks done!')\n",
    "        print('Train Mask2former')\n",
    "        mask2former_id = self.train_mask2former(segformer_id)\n",
    "        print('Mask2former training finished!')\n",
    "\n",
    "        return {'segformer_id': segformer_id, 'mask2former_id': mask2former_id}\n",
    "\n",
    "    def train_segformer(self):\n",
    "        wandb_config = utils.init_wandb('segformer.yml')\n",
    "        segformer_id = {'name': wandb.run.name, 'id': wandb.run.id}\n",
    "        trainer = segformer_tm.get_trainer(self.config)\n",
    "        lightning = segformer_tm.get_lightning(self.config, wandb_config)\n",
    "        trainer.fit(model=lightning)\n",
    "        wandb.finish()\n",
    "\n",
    "        # Deleting the instance to free up memory space.\n",
    "        del trainer, lightning\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        return segformer_id\n",
    "\n",
    "    def train_mask2former(self, segformer_id):\n",
    "        utils.init_wandb('mask2former.yml')\n",
    "        wandb.config.update(dict(model_mask_id=f'{segformer_id[\"name\"]}-{segformer_id[\"id\"]}'))\n",
    "        wandb_config = wandb.config\n",
    "        mask2former_id = {'name': wandb.run.name, 'id': wandb.run.id}\n",
    "        trainer = mask2former_tm.get_trainer(self.config)\n",
    "        lightning = mask2former_tm.get_lightning(self.config, wandb_config)\n",
    "        trainer.fit(model=lightning)\n",
    "        wandb.finish()\n",
    "\n",
    "        # Deleting the instance to free up memory space.\n",
    "        del trainer, lightning\n",
    "        torch.cuda.empty_cache()\n",
    "\n",
    "        return mask2former_id\n",
    "\n",
    "    def predict(self, segformer_id, mask2former_id):\n",
    "        print('Create Segment Anything binary masks...')\n",
    "        self.make_sam_inference(split='test')\n",
    "        print('Segment Anything binary masks done!')\n",
    "        print('Create Segformer binary masks...')\n",
    "        self.make_segformer_inference(split='test', segformer_id=segformer_id)\n",
    "        print('Segformer binary masks done!')\n",
    "        print('Create Mask2former instance masks...')\n",
    "        self.make_mask2former_inference(split='test', mask2former_id=mask2former_id, segformer_id=segformer_id)\n",
    "        print('Mask2former instance masks done!')\n",
    "        \n",
    "        print(f'\\nResults are saved in data/processed/{mask2former_id[\"name\"]}-{mask2former_id[\"id\"]}/test.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 🔮 Predict pipeline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The prediction is done on the volumes located in `data/raw/test` and the final results are saved in `data/processed/vermilion-moon-1241-nvbtr9k2/test`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch size: Min 2 - Max 300\n",
    "segmentation_pipeline = SegmentationPipeline(sam_batch=5, segformer_batch=50, mask2former_batch=25)\n",
    "\n",
    "# To predict with our solution\n",
    "segmentation_pipeline.predict(\n",
    "    segformer_id=dict(id='mmw4795a', name='abundant-lantern-1231'),\n",
    "    mask2former_id=dict(id='nvbtr9k2', name='vermilion-moon-1241'),\n",
    ")\n",
    "\n",
    "# Deleting the instance to free up memory space.\n",
    "del segmentation_pipeline\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ⚙️ Full pipeline (training + infering)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# batch size: Min 2 - Max 300\n",
    "segmentation_pipeline = SegmentationPipeline(sam_batch=5, segformer_batch=50, mask2former_batch=25)\n",
    "\n",
    "# To train and predict new models\n",
    "models_id = segmentation_pipeline.train()\n",
    "segmentation_pipeline.predict(**models_id)\n",
    "\n",
    "# Deleting the instance to free up memory space.\n",
    "del segmentation_pipeline\n",
    "torch.cuda.empty_cache()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}