Duplicate from veb-101/UWMGI_Medical_Image_Segmentation

Co-authored-by: Vaibhav Singh <[email protected]>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+artifacts
+wandb

README.md

@@ -0,0 +1,27 @@
+---
+title: Medical Image Segmentation Gradio App
+emoji: 🏥🩺
+colorFrom: gray
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.37.0
+app_file: app.py
+pinned: false
+duplicated_from: veb-101/UWMGI_Medical_Image_Segmentation
+---
+
+# Medical Image Segmentation Gradio App
+
+For the Gradio app we've removed the dependency on pytorch-lightning otherwise used in the project.
+The `load_lightning_SD_to_Usual_SD.ipynb` notebook contains the steps used to convert pytorch-lightning checkpoint to a regular model checkpoint. This was mainly done to reduce the file size (977 MB --> 244 MB).
+
+You can download the original saved checkpoint from over here: [wandb artifact](https://wandb.ai/veb-101/UM_medical_segmentation/artifacts/model/model-jsr2fn8v/v0/files)
+
+Or via Python:
+
+```python
+import wandb
+run = wandb.init()
+artifact = run.use_artifact('veb-101/UM_medical_segmentation/model-jsr2fn8v:v0', type='model')
+artifact_dir = artifact.download()
+```

app.py

@@ -0,0 +1,77 @@
+import os
+import numpy as np
+import gradio as gr
+from glob import glob
+from functools import partial
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as TF
+from transformers import SegformerForSemanticSegmentation
+
+
+@dataclass
+class Configs:
+    NUM_CLASSES: int = 4  # including background.
+    CLASSES: tuple = ("Large bowel", "Small bowel", "Stomach")
+    IMAGE_SIZE: tuple[int, int] = (288, 288)  # W, H
+    MEAN: tuple = (0.485, 0.456, 0.406)
+    STD: tuple = (0.229, 0.224, 0.225)
+    MODEL_PATH: str = os.path.join(os.getcwd(), "segformer_trained_weights")
+
+
+def get_model(*, model_path, num_classes):
+    model = SegformerForSemanticSegmentation.from_pretrained(model_path, num_labels=num_classes, ignore_mismatched_sizes=True)
+    return model
+
+
+@torch.inference_mode()
+def predict(input_image, model=None, preprocess_fn=None, device="cpu"):
+    shape_H_W = input_image.size[::-1]
+    input_tensor = preprocess_fn(input_image)
+    input_tensor = input_tensor.unsqueeze(0).to(device)
+
+    # Generate predictions
+    outputs = model(pixel_values=input_tensor.to(device), return_dict=True)
+    predictions = F.interpolate(outputs["logits"], size=shape_H_W, mode="bilinear", align_corners=False)
+
+    preds_argmax = predictions.argmax(dim=1).cpu().squeeze().numpy()
+
+    seg_info = [(preds_argmax == idx, class_name) for idx, class_name in enumerate(Configs.CLASSES, 1)]
+
+    return (input_image, seg_info)
+
+
+if __name__ == "__main__":
+    class2hexcolor = {"Stomach": "#007fff", "Small bowel": "#009A17", "Large bowel": "#FF0000"}
+
+    DEVICE = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+
+    model = get_model(model_path=Configs.MODEL_PATH, num_classes=Configs.NUM_CLASSES)
+    model.to(DEVICE)
+    model.eval()
+    _ = model(torch.randn(1, 3, *Configs.IMAGE_SIZE[::-1], device=DEVICE))
+
+    preprocess = TF.Compose(
+        [
+            TF.Resize(size=Configs.IMAGE_SIZE[::-1]),
+            TF.ToTensor(),
+            TF.Normalize(Configs.MEAN, Configs.STD, inplace=True),
+        ]
+    )
+
+    with gr.Blocks(title="Medical Image Segmentation") as demo:
+        gr.Markdown("""<h1><center>Medical Image Segmentation with UW-Madison GI Tract Dataset</center></h1>""")
+        with gr.Row():
+            img_input = gr.Image(type="pil", height=360, width=360, label="Input image")
+            img_output = gr.AnnotatedImage(label="Predictions", height=360, width=360, color_map=class2hexcolor)
+
+        section_btn = gr.Button("Generate Predictions")
+        section_btn.click(partial(predict, model=model, preprocess_fn=preprocess, device=DEVICE), img_input, img_output)
+
+        images_dir = glob(os.path.join(os.getcwd(), "samples") + os.sep + "*.png")
+        examples = [i for i in np.random.choice(images_dir, size=10, replace=False)]
+        gr.Examples(examples=examples, inputs=img_input, outputs=img_output)
+
+    demo.launch()

load_lightning_SD_to_Usual_SD.ipynb

@@ -0,0 +1,335 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Base Configurations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import torch\n",
+    "from transformers import SegformerForSemanticSegmentation\n",
+    "from dataclasses import dataclass\n",
+    "\n",
+    "\n",
+    "@dataclass\n",
+    "class Configs:\n",
+    "    NUM_CLASSES = 4\n",
+    "    MODEL_PATH: str = \"nvidia/segformer-b4-finetuned-ade-512-512\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Load Model To Inspect Parameter Names"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "\n",
+    "def get_model(*, model_path, num_classes):\n",
+    "    model = SegformerForSemanticSegmentation.from_pretrained(\n",
+    "        model_path,\n",
+    "        num_labels=num_classes,\n",
+    "        ignore_mismatched_sizes=True,\n",
+    "    )\n",
+    "    return model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some weights of SegformerForSemanticSegmentation were not initialized from the model checkpoint at nvidia/segformer-b4-finetuned-ade-512-512 and are newly initialized because the shapes did not match:\n",
+      "- decode_head.classifier.weight: found shape torch.Size([150, 768, 1, 1]) in the checkpoint and torch.Size([4, 768, 1, 1]) in the model instantiated\n",
+      "- decode_head.classifier.bias: found shape torch.Size([150]) in the checkpoint and torch.Size([4]) in the model instantiated\n",
+      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "segformer.encoder.patch_embeddings.0.proj.weight\n",
+      "segformer.encoder.patch_embeddings.0.proj.bias\n",
+      "segformer.encoder.patch_embeddings.0.layer_norm.weight\n",
+      "segformer.encoder.patch_embeddings.0.layer_norm.bias\n",
+      "segformer.encoder.patch_embeddings.1.proj.weight\n",
+      "segformer.encoder.patch_embeddings.1.proj.bias\n"
+     ]
+    }
+   ],
+   "source": [
+    "model = get_model(model_path=Configs.MODEL_PATH, num_classes=Configs.NUM_CLASSES)\n",
+    "model_state_dict = model.state_dict()\n",
+    "\n",
+    "print()\n",
+    "for i, (key, val) in enumerate(model_state_dict.items()):\n",
+    "    print(key)\n",
+    "    if i == 5:\n",
+    "        break"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Download & load PyTorch-Lightning Checkpoint and Inspect Parameter Names"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mveb-101\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.15.5"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>c:\\Users\\vaibh\\OneDrive\\Desktop\\Work\\BigVision\\BLOG_POSTS\\Medical_segmentation\\GRADIO_APP\\UWMGI_Medical_Image_Segmentation\\wandb\\run-20230719_204221-w5qu5rqw</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/veb-101/UWMGI_Medical_Image_Segmentation/runs/w5qu5rqw' target=\"_blank\">ethereal-bush-2</a></strong> to <a href='https://wandb.ai/veb-101/UWMGI_Medical_Image_Segmentation' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/veb-101/UWMGI_Medical_Image_Segmentation' target=\"_blank\">https://wandb.ai/veb-101/UWMGI_Medical_Image_Segmentation</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/veb-101/UWMGI_Medical_Image_Segmentation/runs/w5qu5rqw' target=\"_blank\">https://wandb.ai/veb-101/UWMGI_Medical_Image_Segmentation/runs/w5qu5rqw</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Downloading large artifact model-fpgquxev:v0, 977.89MB. 1 files... \n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n",
+      "Done. 0:1:5.3\n"
+     ]
+    }
+   ],
+   "source": [
+    "import wandb\n",
+    "\n",
+    "run = wandb.init()\n",
+    "artifact = run.use_artifact(\"veb-101/UM_medical_segmentation/model-fpgquxev:v0\", type=\"model\")\n",
+    "artifact_dir = artifact.download()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "dict_keys(['epoch', 'global_step', 'pytorch-lightning_version', 'state_dict', 'loops', 'callbacks', 'optimizer_states', 'lr_schedulers', 'MixedPrecisionPlugin', 'hparams_name', 'hyper_parameters'])\n"
+     ]
+    }
+   ],
+   "source": [
+    "CKPT = torch.load(os.path.join(artifact_dir, \"model.ckpt\"), map_location=\"cpu\")\n",
+    "print(CKPT.keys())"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "model.segformer.encoder.patch_embeddings.0.proj.weight\n",
+      "model.segformer.encoder.patch_embeddings.0.proj.bias\n",
+      "model.segformer.encoder.patch_embeddings.0.layer_norm.weight\n",
+      "model.segformer.encoder.patch_embeddings.0.layer_norm.bias\n",
+      "model.segformer.encoder.patch_embeddings.1.proj.weight\n",
+      "model.segformer.encoder.patch_embeddings.1.proj.bias\n"
+     ]
+    }
+   ],
+   "source": [
+    "TRAINED_CKPT_state_dict = CKPT[\"state_dict\"]\n",
+    "\n",
+    "for i, (key, val) in enumerate(TRAINED_CKPT_state_dict.items()):\n",
+    "    print(key)\n",
+    "    if i == 5:\n",
+    "        break"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "**The pytorch-lightning `state_dict()` has an extra `model.` string at the front that refers to the object/variable name that was holding the model in the `LightningModule` class.**\n",
+    "\n",
+    "We can simply iterate over the parameters and change the parameter key name. We'll create a new `OrderedDict` for it."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from collections import OrderedDict\n",
+    "\n",
+    "new_state_dict = OrderedDict()\n",
+    "\n",
+    "for key_name, value in CKPT[\"state_dict\"].items():\n",
+    "    new_state_dict[key_name.replace(\"model.\", \"\")] = value"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<All keys matched successfully>"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.load_state_dict(new_state_dict)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# torch.save(model.state_dict(), \"Segformer_best_state_dict.ckpt\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.save_pretrained(\"segformer_trained_weights\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To load the saved model, we simply need to pass the path to the directory \"segformer_trained_weights\"."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model = get_model(model_path=os.path.join(os.getcwd(), \"segformer_trained_weights\"), num_classes=Configs.NUM_CLASSES)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "pytorchx",
+   "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.12"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

requirements.txt

@@ -0,0 +1,4 @@
+--find-links https://download.pytorch.org/whl/torch_stable.html
+torch==2.0.0+cpu
+torchvision==0.15.0
+transformers==4.30.2

segformer_trained_weights/config.json

@@ -0,0 +1,82 @@
+{
+  "_name_or_path": "nvidia/segformer-b4-finetuned-ade-512-512",
+  "architectures": [
+    "SegformerForSemanticSegmentation"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "classifier_dropout_prob": 0.1,
+  "decoder_hidden_size": 768,
+  "depths": [
+    3,
+    8,
+    27,
+    3
+  ],
+  "downsampling_rates": [
+    1,
+    4,
+    8,
+    16
+  ],
+  "drop_path_rate": 0.1,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_sizes": [
+    64,
+    128,
+    320,
+    512
+  ],
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3"
+  },
+  "image_size": 224,
+  "initializer_range": 0.02,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2,
+    "LABEL_3": 3
+  },
+  "layer_norm_eps": 1e-06,
+  "mlp_ratios": [
+    4,
+    4,
+    4,
+    4
+  ],
+  "model_type": "segformer",
+  "num_attention_heads": [
+    1,
+    2,
+    5,
+    8
+  ],
+  "num_channels": 3,
+  "num_encoder_blocks": 4,
+  "patch_sizes": [
+    7,
+    3,
+    3,
+    3
+  ],
+  "reshape_last_stage": true,
+  "semantic_loss_ignore_index": 255,
+  "sr_ratios": [
+    8,
+    4,
+    2,
+    1
+  ],
+  "strides": [
+    4,
+    2,
+    2,
+    2
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.30.2"
+}

segformer_trained_weights/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:423ff60b52bdbc5c0ea00f1a5648c42eccf2bdfbab550304bc95e28eb594cf0e
+size 256300245