medsam_experimenting_2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Source](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/SAM/Fine_tune_SAM_(segment_anything)_on_a_custom_dataset.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import os, pdb\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import pandas as pd\n",
    "\n",
    "import torch.nn.functional as F\n",
    "from typing import Tuple\n",
    "from torch.nn import MSELoss\n",
    "from transformers import SamModel \n",
    "from transformers import SamProcessor\n",
    "from tqdm import tqdm\n",
    "from statistics import mean\n",
    "import torch\n",
    "from torch.nn.functional import threshold, normalize\n",
    "from torch.optim import Adam\n",
    "import monai\n",
    "from monai.transforms import ToTensor\n",
    "\n",
    "from torchvision.transforms import ToPILImage, Compose, Resize, Grayscale, Lambda\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data import Dataset\n",
    "from datasets import load_dataset  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIR = '/sddata/data/MedSAM/public_test_data_01_10_2023'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"nielsr/breast-cancer\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "example = dataset[0]\n",
    "image = example[\"image\"]\n",
    "type(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = example['label']\n",
    "label = np.array(label) * 255\n",
    "label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(label)[60:180,30:90]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(label, cmap='gray')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "        # color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    print(h, w)\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "\n",
    "fig, axes = plt.subplots()\n",
    "\n",
    "axes.imshow(np.array(image))\n",
    "ground_truth_seg = np.array(example[\"label\"])\n",
    "show_mask(ground_truth_seg, axes)\n",
    "axes.title.set_text(f\"Ground truth mask\")\n",
    "axes.axis(\"off\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Pytorch Dataset\n",
    "\n",
    "Below we define a regular PyTorch dataset, which gives us examples of the data prepared in the format for the model. Each example consists of:\n",
    "\n",
    "* pixel values (which is the image prepared for the model)\n",
    "* a prompt in the form of a bounding box\n",
    "* a ground truth segmentation mask.\n",
    "The function below defines how to get a bounding box prompt based on the ground truth segmentation. This was taken from here.\n",
    "\n",
    "Note that SAM is always trained using certain \"prompts\", which you could be bounding boxes, points, text, or rudimentary masks. The model is then trained to output the appropriate mask given the image + prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_bounding_box(ground_truth_map):\n",
    "  # get bounding box from mask\n",
    "  y_indices, x_indices = np.where(ground_truth_map > 0)\n",
    "  x_min, x_max = np.min(x_indices), np.max(x_indices)\n",
    "  y_min, y_max = np.min(y_indices), np.max(y_indices)\n",
    "  # add perturbation to bounding box coordinates\n",
    "  H, W = ground_truth_map.shape\n",
    "  x_min = max(0, x_min - np.random.randint(0, 20))\n",
    "  x_max = min(W, x_max + np.random.randint(0, 20))\n",
    "  y_min = max(0, y_min - np.random.randint(0, 20))\n",
    "  y_max = min(H, y_max + np.random.randint(0, 20))\n",
    "  bbox = [x_min, y_min, x_max, y_max]\n",
    "\n",
    "  return bbox"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAMDataset(Dataset):\n",
    "    def __init__(self, dataset, processor):\n",
    "        self.dataset = dataset\n",
    "        self.processor = processor\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        item = self.dataset[idx]\n",
    "        image = item[\"image\"]\n",
    "        ground_truth_mask = np.array(item[\"label\"])\n",
    "\n",
    "        # get bounding box prompt\n",
    "        prompt = get_bounding_box(ground_truth_mask)\n",
    "\n",
    "        # prepare image and prompt for the model\n",
    "        inputs = self.processor(image, input_boxes=[[prompt]], return_tensors=\"pt\")\n",
    "\n",
    "        # remove batch dimension which the processor adds by default\n",
    "        inputs = {k:v.squeeze(0) for k,v in inputs.items()}\n",
    "\n",
    "        # add ground truth segmentation\n",
    "        inputs[\"ground_truth_mask\"] = ground_truth_mask\n",
    "\n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = SamProcessor.from_pretrained(\"facebook/sam-vit-base\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = SAMDataset(dataset=dataset, processor=processor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "example = train_dataset[0]\n",
    "for k,v in example.items():\n",
    "    print(k,v.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create PyTorch DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataloader = DataLoader(train_dataset, batch_size=10, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = next(iter(train_dataloader))\n",
    "for k,v in batch.items():\n",
    "    print(k,v.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "for item in train_dataloader:\n",
    "    count += 1\n",
    "    print(count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = SamModel.from_pretrained(\"facebook/sam-vit-base\")\n",
    "\n",
    "# make sure we only compute gradients for mask decoder\n",
    "for name, param in model.named_parameters():\n",
    "    if name.startswith(\"vision_encoder\") or name.startswith(\"prompt_encoder\"):\n",
    "        param.requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Note: Hyperparameter tuning could improve performance here\n",
    "optimizer = Adam(model.mask_decoder.parameters(), lr=1e-5, weight_decay=0)\n",
    "\n",
    "seg_loss = monai.losses.DiceCELoss(sigmoid=True, squared_pred=True, reduction='mean')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_epochs = 5\n",
    "\n",
    "# device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device = \"cpu\"\n",
    "model.to(device)\n",
    "\n",
    "model.train()\n",
    "epoch_losses = []\n",
    "for epoch in range(num_epochs):\n",
    "    batch_epoch_losses = []\n",
    "    for batch in tqdm(train_dataloader):\n",
    "        # forward pass\n",
    "        outputs = model(pixel_values=batch[\"pixel_values\"].to(device),\n",
    "                        input_boxes=batch[\"input_boxes\"].to(device),\n",
    "                        multimask_output=False)\n",
    "        \n",
    "        # compute loss\n",
    "        pdb.set_trace()\n",
    "        predicted_masks = outputs.pred_masks.squeeze(1)\n",
    "        ground_truth_masks = batch[\"ground_truth_mask\"].float().to(device)\n",
    "        loss = seg_loss(predicted_masks, ground_truth_masks.unsqueeze(1))\n",
    "        \n",
    "        # backward pass (compute gradients of parameters w.r.t. loss)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        \n",
    "        # optimize\n",
    "        optimizer.step()\n",
    "        batch_epoch_losses.append(loss.item())\n",
    "\n",
    "    print(f'EPOCH: {epoch}')\n",
    "    print(f'Mean loss: {mean(batch_epoch_losses)}')\n",
    "    epoch_losses.append(mean(batch_epoch_losses))\n",
    "    plt.figure()\n",
    "    plt.xlabel('epoch')\n",
    "    plt.ylabel('loss')\n",
    "    plt.plot(epoch_losses)\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save it"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the trained model\n",
    "torch.save(model.state_dict(), \"/sddata/data/MedSAM/model/facebook_sam-vit-base_BB.pth\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use our own data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Image\n",
    "raw_image = Image.open('/sddata/data/MedSAM/public_test_data_01_10_2023/Color_Fundus/visivite_GA.png')\n",
    "# Convert RGBA image to RGB\n",
    "raw_image = raw_image.convert(\"RGB\")\n",
    "\n",
    "# Label\n",
    "label = Image.open('/sddata/data/MedSAM/public_test_data_01_10_2023/labels/visivite_GA.png')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(type(raw_image))\n",
    "print(type(label))\n",
    "print(np.array(raw_image).shape)\n",
    "print(np.array(label).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print((np.array(label)[:,:,0] == 0).all()) # R\n",
    "print((np.array(label)[:,:,1] == 0).all()) # G\n",
    "print((np.array(label)[:,:,2] == 0).all()) # B\n",
    "print((np.array(label)[:,:,3] == 0).all()) # A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.array(label)[800:900,800:900,0] # R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "w = np.array(label).shape[0]\n",
    "h = np.array(label).shape[1]\n",
    "ground_truth_seg = np.array(label)[:,:,0].reshape(w, h)\n",
    "ground_truth_seg_norm = ground_truth_seg / 255\n",
    "print(f\"{w} {h}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ground_truth_seg_norm.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots()\n",
    "\n",
    "axes.imshow(np.array(raw_image))\n",
    "show_mask(ground_truth_seg_norm, axes, random_color=True)\n",
    "axes.title.set_text(f\"Ground truth mask\")\n",
    "axes.axis(\"off\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SAMDataset(Dataset):\n",
    "    def __init__(self, dataset, processor, is_grayscale=False):\n",
    "        if isinstance(dataset, str):\n",
    "            self.dataset = pd.read_csv(os.path.join(dataset))\n",
    "        else:\n",
    "            self.dataset = dataset\n",
    "        self.processor = processor\n",
    "        self.is_grayscale = is_grayscale\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "    def _load_image(self, image_path, is_grayscale=True, resize=(256, 256)):\n",
    "        image = Image.open(os.path.join(DATA_DIR,image_path))\n",
    "        image = image.convert('RGB')\n",
    "        # image = image.convert('RGB')\n",
    "        img_transform = Compose([\n",
    "            Resize(resize),\n",
    "            Grayscale() if is_grayscale else Lambda(lambda x: x),\n",
    "            # ToTensor() # it will convert to tensor and scale between 0-1\n",
    "            Lambda(lambda x: torch.from_numpy(np.array(x))),\n",
    "            \n",
    "        ])\n",
    "        image = img_transform(image)\n",
    "        # pdb.set_trace()\n",
    "        return image\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if isinstance(self.dataset, pd.DataFrame):\n",
    "            item = self.dataset.iloc[idx]\n",
    "            image = item['image_path_orig']\n",
    "            label = item['labels']\n",
    "            image = self._load_image(image, is_grayscale=False, resize=(256, 256))\n",
    "            label = self._load_image(label, is_grayscale=True, resize=(256, 256))\n",
    "            label = (label > 0).float()\n",
    "        else:            \n",
    "            item = self.dataset[idx]\n",
    "            image = item[\"image\"]\n",
    "            label = np.array(item[\"label\"])\n",
    "\n",
    "        # get bounding box prompt\n",
    "        # print(image)\n",
    "        # print(type(image))\n",
    "        # print(label.shape)\n",
    "        prompt = get_bounding_box(label.squeeze())\n",
    "\n",
    "        # prepare image and prompt for the model\n",
    "        print(image.min(), image.max())\n",
    "        inputs = self.processor(image, input_boxes=[[prompt]], return_tensors=\"pt\")\n",
    "        # print(image.min(), image.max())\n",
    "        # pdb.set_trace()\n",
    "        # remove batch dimension which the processor adds by default\n",
    "        inputs = {k:v.squeeze(0) for k,v in inputs.items()}\n",
    "        # inputs = {'image': image, 'label': label}\n",
    "\n",
    "        # add ground truth segmentation\n",
    "        inputs[\"ground_truth_mask\"] = label\n",
    "\n",
    "        return inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Pytorch Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = os.path.join(DATA_DIR, \"image_key_cf.csv\")\n",
    "dataset_ours = SAMDataset(data_path, processor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(3, dtype=torch.uint8) tensor(255, dtype=torch.uint8)\n",
      "tensor(3, dtype=torch.uint8) tensor(255, dtype=torch.uint8)\n"
     ]
    }
   ],
   "source": [
    "index = 0\n",
    "image = dataset_ours[index]['pixel_values']\n",
    "label = dataset_ours[index]['ground_truth_mask']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "/9j/4AAQSkZJRgABAQAAAQABAAD/2wBDAAgGBgcGBQgHBwcJCQgKDBQNDAsLDBkSEw8UHRofHh0aHBwgJC4nICIsIxwcKDcpLDAxNDQ0Hyc5PTgyPC4zNDL/wAALCAEAAQABAREA/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/9oACAEBAAA/APn+iiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiiipIIJrq4it7eKSaeVwkccalmdicAADkknjFeueDfgDruteVeeIZf7IsWw3k4DXMi/KcbekeQWGWyykcpXt+jfC3wVomnJZxeHrG6xgtNfQrcSO2ACSzg4zjOFwuScAZrQ/wCEE8H/APQqaH/4Lof/AImuf/4Ul8PP+he/8nbj/wCOUf8ACkvh5/0L3/k7cf8AxyvJPif8EpvDdvda94daS50tHLzWZUmS0jwDkNkl0Bzk9VGM7sMw8booooooooooooooooooooooor1j4f8AwO1TxXZ22r6tc/2ZpM2HjULmedMjlQeFVhuwxz0B2kEGvoPwp4D8OeDLcJo2nRxzlNsl3J888nC5y55AJUHaMLnkAV0lFFFFFfJHxe+G/wDwg2uJc6bDOdCvOYXf5hBJzmEtnJwBkFsEjj5irGvN6KKKKKKKKKKKKKKKKKKKKK98+GHwMm+0Wuu+L4YxAEEsGluCWLZOPPBGAAMHZznOGxgqfoOiiiiiiiqeq6VY65pdxpmp20dzZ3CbJYn6MP5gg4II5BAIwRXyp8SPhDqngbzdStn+3aEZdqTD/WQA42iUYwMk7Qw4JAztLBa83ooooooooooooooooooor6H+CPwst47O08Y63F5txJ+80+1kQgRAHiZgRyxxlewBDckjb7xRRRRRRRRRRXzx8U/gjcR3kmt+DrLzbeTc9zpsWAYiASWiHdTj7g5BICgg4XweiiiiiiiiiiiiiiiiivSPhD8N/wDhOdce51KGcaFZ8zOnyieTjEIbORkHJK5IHHyllNfW9FFFFFFFFFFFFeN/G34YTeJLdfEWg2kb6pboRdwxqfMu4wBgjsXUA8YywOMnaqn5kooooooooooooooorU8N6K3iLxLpujpPHAb24SHzXKgICeT8xAJx0XOWOAOSK+1/DmgWPhbw/Z6LpqyC0tUKp5jbmYklmYn1LEnjA54AHFalFFFFFFFFFFFFFfCGu6Z/YniHU9J87zvsN3Lbebt279jld2MnGcZxk1n0UUUUUUUUUUUUUVJBPNa3EVxbyyQzxOHjkjYqyMDkEEcgg85r7b8D+JP+Eu8F6Xrhj8uS6i/eoFwBIpKPtGT8u5Wxk5xjPNdBRRRRRRRRRRRRRXyR8dbG4tPivqU08eyO7igmgO4HegjWMnjp8yMOfT0xXm9FFFFFFFFFFFFFFFfX/wAEv+SQ6F/28f8ApRJXoFFFFFFFFFFFFFFfOn7S2mwxazoGqK0nn3FvLbupI2hY2VlI4znMrZ57Dp38Looooooooooooooor7X+HFjb6f8ADXw5Dax+XG2nwzEbicvIokc8+rMx9s8cV1FFFFFFFFFFFFFFeP8A7Rlh9o8A2V4lp5klrqCbphHkxRsjg5b+FS3lg9idvfFfMFFFFFFFFFFFFFFFFfY/wj8Q2/iH4a6S8C7JLCJbCdMk7XiUAckDOV2txnG7GSQa7iiiiiiiiiiiiiivnD43fFL+1JbvwbpC4s4ZfL1Cd05lkRs+WoPRVZRlupI4+UZbw+iiiiiiiiiiiiiiivUPgX4u/wCEc8cpptw2LHWdts/H3ZgT5TcKSeSUxkD95k/dr6voooooooooooory/41+Pf+EU8LnTdOvfJ1vUMLH5T4kghyd0nQ4zjYDkHLEqcoa+UKKKKKKKKKKKKKKKKKK+v/AIPeMv8AhL/A1v8AaZvM1PT8Wt3ubLvgfJIcsWO5erHGWV8dK9Aoooooooooorl/HXjrS/Aehm/vz5txJlbW0RsPcOOw9FGRlu2e5IB+QPE/ifVPF2uTavq8/m3EnCqvCRIOiIOyjJ/MkkkknHooooooooooooooooor0D4Q+OrfwN4tebUTP/Zl7F5Fx5bEiM7gVlKD723keoDtjPQ/W9jf2ep2cd5YXcF3ayZ2TQSCRGwSDhhwcEEfhViiiiiiiiiuP+JXjf8A4QLwk2qx2v2m6llFtbIxwgkZWYM/OdoCk4HJ4HGcj5I8T+J9U8Xa5Nq+rz+bcScKq8JEg6Ig7KMn8ySSSSceiiiiiiiiiiiiiiiiiiiiivUPCPx08U+HNtvqT/23YjPyXchEy/ePEuCTyR94NwoAxXv/AIR+J3hbxpti02+8q+Of9BuwI5v4ugyQ/ClvlLYGM4rsKKKy9b8SaL4ct/P1nVLSxQo7oJpQrSBRlti9XIyOFBPI9a5P/hdvw8/6GH/ySuP/AI3UmpfGXwHptusp12O5d7cXEcVrE8jMCDhTgYRzjG1ypGRnFeAfFn4kw/ELVLEWNnJb6fYI4hM+PNkZ9u8sASAPlAAGehJPOB53RRRRRRRRRRRRRRRRRRRRRRRUkE81rcRXFvLJDPE4eOSNirIwOQQRyCDzmve/Av7QaRWYsvGazy3HmgJqFtCuCjHkyIMY2+qA5HbIy3pF58XvA1poY1Ya9BcRtuEdvCCZ3YZ48o4Zc7SAWCryOcEV4h46+Omu+IpTa6A8+i6aMgmOQefN82VYuBmPgD5VPdgWYHjyueea6uJbi4lkmnlcvJJIxZnYnJJJ5JJ5zUdFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFf/9k=",
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAQAAAAEACAAAAAB5Gfe6AAAEXUlEQVR4Ae3c7W7bMAyFYXfY/d9y5qEr4g85IykeSjbe/VnsJiLPI9pJi27Lwh8EEEAAAQQQQAABBBBAAAEEEEAAAQQQQAABBBBAAAEEEEAAAQQQQAABBBBAAAEEEEAAAQQQQAABBBBAAAEEEEAAAQQQQAABBBBAAAEEEEAAAQQQQAABBBBAAAEEEEAAAQQQQOBmAl/Cfl/btZWFtnW8j3V97eI32tJVbhS7PqVq43/x145Upa/Dtr6i6cIQ/28zmuKtmNfnFD0Y488hkA9gjz+FQDaAK/4MArkA3vgTCPxee8j6E4ifVTq+Tt4EROPndRBSyCofjb82ndVCKH9S9Y74owUy7gF98WMbl/aqX/0rdefvXqAnQz9AQvsJS4QNugFSmk9ZJGbQC5DUetIyAYNOgLTG0xbyGvQBJLaduJTLoAsgtenUxewGHR/D8jvuaMae+PDM+ATk518ESx7ing/DAJJmJYueQ2/PRKdO1mq0oW0oz+PgBMjyl18GMQBd/qVaIASgzF8tEAHQ5i8WCACo89cKBAA899jgc1965J/O/AA1vdVUWRXcb7tlnflb+9lU19/+CXAt3/PkGmovQE1X324ltZwAJT31jI37tT6A4vwV5XwAbt/OFxQIuAAK+jmI6St6APTdHOKvh/KaDgB5L+f4BWccAAXdDChhB3jmACxmgIfmNwMMy68ubJ6AAZdnSUkjgHofSrI2ixgB3N81N4vNeNIIMLB18fBZAR47AlaAgSOgLW0GeOoImAG0+zBudTvAQ0fADjBsk7RvAzcA0MoDoPWdf3UmYP490v5ckAm4wQRIW2QCpLw3WJwJuMEmSd8G7BOg/Ug+bB/sAMNa1BYGQOs7/+pMwPx7tHYovAEzAbeYAGGT95gA4Q9kbwEgzG/+/YDAbxWnza0yvwMgLY93IWl+D4C2kUsWcdlb3AMucRK+4AEQ70U7jbqoB6DdofasOr/nHjDijUCe3weg3ewxq/suAf2G7BUK6vkA9u094sgJULAlG9aKak6ATXf6hxX53TfBkqb0tu8K7gmoE6ip5AZ424kf1eR3XwJln4aK8gcAxBv/b/mq/BGAit4qanxTz3kPqMsfmQD9v2svzB8CUN8FKvPHALQdalc/7t5894Da/EvG/yp7RO05Lo6/xC4B3W2wPH8QQCVQnz8KoBLouXpirw3fBL8G7FYs4udXpcXI+B2GtGY+Z959NTwBu1XWg7tORBrAXQmSp67rQkju5Tij7eP0onGC9FbaiQ9nBVVDBII+DkkvDiWFDQTvuq+xnynejVwIxU5/JhAVDbWq6uVaQFUxFF/4M962wGTpVzRhR2cCYbHg/ke/HTaVO8U9nTAtI36StKndDEgrxZXUbf0YqOuEBQoaWw0KqoQFeCECCCCAAAIIIIAAAggggAACCCCAAAIIIIAAAggggAACCCCAAAIIIIAAAggggAACCCCAAAIIIIAAAggggAACCCCAAAIIIPAUgT+t6lTTARnGtgAAAABJRU5ErkJggg==",
      "text/plain": [
       "<PIL.Image.Image image mode=L size=256x256>"
      ]
     },
     "execution_count": 99,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ToPILImage()(image)\n",
    "ToPILImage()(label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 1024, 1024])\n",
      "torch.Size([256, 256])\n"
     ]
    }
   ],
   "source": [
    "print(image.shape)\n",
    "print(label.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_ours_dataloader = DataLoader(dataset_ours, batch_size=3, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0, dtype=torch.uint8) tensor(255, dtype=torch.uint8)\n",
      "tensor(3, dtype=torch.uint8) tensor(255, dtype=torch.uint8)\n",
      "tensor(0, dtype=torch.uint8) tensor(251, dtype=torch.uint8)\n",
      "> \u001b[0;32m/tmp/ipykernel_103234/3505833551.py\u001b[0m(26)\u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;32m     24 \u001b[0;31m        \u001b[0;31m# compute loss\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m     25 \u001b[0;31m        \u001b[0mpdb\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mset_trace\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m---> 26 \u001b[0;31m        \u001b[0mpredicted_masks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0moutputs\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpred_masks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m     27 \u001b[0;31m        \u001b[0mground_truth_masks\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mbatch\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"ground_truth_mask\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfloat\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\u001b[0;32m     28 \u001b[0;31m        \u001b[0mloss\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mseg_loss\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpredicted_masks\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mground_truth_masks\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munsqueeze\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0m\n",
      "tensor([0., 1.])\n",
      "tensor([0., 1.])\n",
      "torch.Size([3, 256, 256])\n",
      "tensor([-10.1430, -10.1292,  -9.8847,  ...,   8.0285,   8.0409,   8.2014],\n",
      "       grad_fn=<Unique2Backward0>)\n",
      "tensor(0.2498, grad_fn=<AddBackward0>)\n",
      "torch.Size([])\n",
      "tensor(0.2498, grad_fn=<AddBackward0>)\n",
      "tensor([[[[-7.0220, -7.3053, -6.5896,  ..., -7.1348, -6.7828, -6.9981],\n",
      "          [-7.1682, -7.4578, -6.9291,  ..., -7.2889, -7.2929, -7.2047],\n",
      "          [-6.8061, -7.1369, -6.9731,  ..., -7.1041, -7.0495, -7.0763],\n",
      "          ...,\n",
      "          [-6.9424, -7.4183, -6.9559,  ..., -7.2318, -7.2576, -7.0051],\n",
      "          [-6.8940, -7.3817, -6.8126,  ..., -7.1235, -6.6851, -6.6310],\n",
      "          [-6.7942, -7.0091, -6.9211,  ..., -6.7099, -6.7599, -6.7906]]],\n",
      "\n",
      "\n",
      "        [[[-3.9119, -4.2128, -3.7011,  ..., -4.1664, -3.9083, -4.1022],\n",
      "          [-4.0571, -4.0249, -3.9755,  ..., -3.9579, -3.9426, -3.8401],\n",
      "          [-3.7346, -3.9822, -3.7888,  ..., -4.3912, -4.2632, -4.2393],\n",
      "          ...,\n",
      "          [-4.4927, -5.0329, -4.7165,  ..., -4.8793, -4.6265, -4.3381],\n",
      "          [-3.9985, -4.2082, -3.6455,  ..., -3.9551, -3.4367, -3.5960],\n",
      "          [-4.2019, -4.6238, -4.0343,  ..., -4.1256, -3.6328, -3.8473]]],\n",
      "\n",
      "\n",
      "        [[[-3.3253, -3.4557, -3.0138,  ..., -3.2229, -2.8813, -3.1149],\n",
      "          [-3.2907, -3.2895, -3.0410,  ..., -3.0135, -2.9558, -2.9906],\n",
      "          [-3.2459, -3.3668, -3.2048,  ..., -3.1834, -3.0572, -3.0662],\n",
      "          ...,\n",
      "          [-3.3652, -3.5215, -3.2653,  ..., -3.2421, -2.9624, -3.2409],\n",
      "          [-3.4518, -3.4621, -3.1311,  ..., -3.1138, -2.7396, -2.9025],\n",
      "          [-3.2244, -3.4271, -3.2617,  ..., -3.1376, -2.8724, -3.1763]]]],\n",
      "       grad_fn=<SqueezeBackward1>)\n",
      "tensor([[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         ...,\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "        [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         ...,\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.]],\n",
      "\n",
      "        [[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         ...,\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.]]])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/1 [23:04<?, ?it/s]\n"
     ]
    }
   ],
   "source": [
    "model = SamModel.from_pretrained(\"facebook/sam-vit-base\")\n",
    "\n",
    "# make sure we only compute gradients for mask decoder\n",
    "for name, param in model.named_parameters():\n",
    "    if name.startswith(\"vision_encoder\") or name.startswith(\"prompt_encoder\"):\n",
    "        param.requires_grad_(False)\n",
    "\n",
    "num_epochs = 15\n",
    "\n",
    "# device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "device = \"cpu\"\n",
    "model.to(device)\n",
    "\n",
    "model.train()\n",
    "epoch_losses = []\n",
    "for epoch in range(num_epochs):\n",
    "    batch_epoch_losses = []\n",
    "    for batch in tqdm(train_dataset_ours_dataloader):\n",
    "        # forward pass\n",
    "        outputs = model(pixel_values=batch[\"pixel_values\"].to(device),\n",
    "                        input_boxes=batch[\"input_boxes\"].to(device),\n",
    "                        multimask_output=False)\n",
    "        \n",
    "        # compute loss\n",
    "        pdb.set_trace()\n",
    "        predicted_masks = outputs.pred_masks.squeeze(1)\n",
    "        ground_truth_masks = batch[\"ground_truth_mask\"].float().to(device)\n",
    "        loss = seg_loss(predicted_masks, ground_truth_masks.unsqueeze(1))\n",
    "        \n",
    "        # backward pass (compute gradients of parameters w.r.t. loss)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        \n",
    "        # optimize\n",
    "        optimizer.step()\n",
    "        batch_epoch_losses.append(loss.item())\n",
    "\n",
    "    print(f'EPOCH: {epoch}')\n",
    "    print(f'Mean loss: {mean(batch_epoch_losses)}')\n",
    "    epoch_losses.append(mean(batch_epoch_losses))\n",
    "    plt.figure()\n",
    "    plt.xlabel('epoch')\n",
    "    plt.ylabel('loss')\n",
    "    plt.plot(epoch_losses)\n",
    "    plt.show()\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Pytorch DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train The Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "MedSAM",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}