notebooks/visualize_hierarchical_segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28bda3bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import sys\n",
    "from functools import partial\n",
    "\n",
    "import cv2\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.colors as colors\n",
    "import skimage.color as sk_color\n",
    "import skimage.morphology as sk_morph\n",
    "import scipy.io\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "sys.path.append('../moco-v3/')\n",
    "import suppix_utils.datasets_seeds as datasets\n",
    "\n",
    "sys.path.append('../')\n",
    "import cast_models.cast"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "355485a3",
   "metadata": {},
   "source": [
    "### Prepare utility functions for visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cbafe93",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_contour(label, size=2, is_symmetric=True):\n",
    "    if is_symmetric:\n",
    "        label = np.pad(label, ((1, 1), (1, 1)), mode='symmetric')\n",
    "    else:\n",
    "        label = np.pad(label, ((1, 1), (1, 1)), mode='constant', constant_values=0)\n",
    "    l_diff = label[1:-1, 1:-1] != label[1:-1, :-2]\n",
    "    r_diff = label[1:-1, 1:-1] != label[1:-1, 2:]\n",
    "    t_diff = label[1:-1, 1:-1] != label[:-2, 1:-1]\n",
    "    b_diff = label[1:-1, 1:-1] != label[2:, 1:-1]\n",
    "\n",
    "    edge = (l_diff + r_diff + t_diff + b_diff).astype(label.dtype)\n",
    "\n",
    "    # dilation\n",
    "    if size > 0:\n",
    "        disk = sk_morph.disk(size)\n",
    "        edge = edge.astype(np.int32)\n",
    "        edge = sk_morph.dilation(edge, disk).astype(label.dtype)\n",
    "    return edge\n",
    "\n",
    "def label2color(label, img):\n",
    "    out = sk_color.label2rgb(label, img, kind='avg', bg_label=-1)\n",
    "    edge = generate_contour(label, 0).astype(bool)\n",
    "    out[edge] = 1\n",
    "    return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f03356b",
   "metadata": {},
   "source": [
    "### Load color mappings for visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d89060",
   "metadata": {},
   "outputs": [],
   "source": [
    "CMAP = scipy.io.loadmat('colormapvoc.mat')['colormapvoc']\n",
    "CMAP = (CMAP * 255).astype(np.uint8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40b2e17f",
   "metadata": {},
   "source": [
    "### Define models and load pre-trained weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e279f7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MODEL_CLASS_NAME = 'cast_base'\n",
    "CHECKPOINT_PATH = '../snapshots/moco/imagenet1k/cast_base/checkpoint_0099.pth.tar'\n",
    "\n",
    "model = cast_models.cast.__dict__[MODEL_CLASS_NAME]().cuda()\n",
    "ckpt = torch.load(CHECKPOINT_PATH, map_location='cuda:0')\n",
    "state_dict = {k[len('module.base_encoder.'):]: v for k, v in ckpt['state_dict'].items()\n",
    "              if 'module.base_encoder.' in k and 'head' not in k}\n",
    "msg = model.load_state_dict(state_dict, strict=False)\n",
    "print(msg)\n",
    "model = model.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b62f3b26",
   "metadata": {},
   "source": [
    "### Prepare dataloader for loading ImageNet images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "373ceba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "normalize = transforms.Normalize(\n",
    "    mean=[0.485, 0.456, 0.406],\n",
    "    std=[0.229, 0.224, 0.225])\n",
    "\n",
    "augmentation = [\n",
    "    transforms.Resize(224),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "]\n",
    "\n",
    "train_dataset = datasets.ImageFolder(\n",
    "    'demo_images',\n",
    "    transforms.Compose(augmentation),\n",
    "    normalize=normalize,\n",
    "    n_segments=196,\n",
    "    compactness=10.0,\n",
    "    blur_ops=None,\n",
    "    scale_factor=1.0)\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    train_dataset, batch_size=1, shuffle=False, drop_last=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffa9c2ca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "for i, (img, suppixel, label) in enumerate(train_loader):\n",
    "    img = img.cuda() # input images\n",
    "    suppixel = suppixel.cuda() # superpixels\n",
    "\n",
    "    # Forward pass to return intermediate groupings\n",
    "    intermediates = model.forward_features(img, suppixel)\n",
    "\n",
    "    # Aggregate groupings from fine to coarse levels\n",
    "    segmentations = {}\n",
    "    prev_labels = {}\n",
    "    # Iterate through the finest to the coarsest scales\n",
    "    for level in [1, 2, 3, 4]:\n",
    "        # Iterate through the mini-batch\n",
    "        for b in range(img.shape[0]):\n",
    "            # Grouping logit for the current level\n",
    "            logit = intermediates['logit{:d}'.format(level)][b]\n",
    "            label = torch.argmax(logit, dim=-1)\n",
    "\n",
    "            # Gather coarser grouping at the current level\n",
    "            # The level-1 grouping index for each level-0 group is [1, 2, 2, 0, 1, 1, 0]\n",
    "            # The level-2 grouping index for each level-1 group is [0, 1, 0]\n",
    "            # We infer the level-2 grouping for each level-0 group as [1, 0, 0, 0, 1, 1, 0]\n",
    "            if level > 1:\n",
    "                prev_label = prev_labels['level{:d}'.format(level-1)][b]\n",
    "                label = torch.gather(label, 0, prev_label.view(-1))\n",
    "            if prev_labels.get('level{:d}'.format(level), None) is None:\n",
    "                prev_labels['level{:d}'.format(level)] = []\n",
    "            prev_labels['level{:d}'.format(level)].append(label)\n",
    "\n",
    "            # Gather groupings for each superpixel\n",
    "            label = torch.gather(label, 0, suppixel[b].view(-1))\n",
    "            label = label.view(suppixel.shape[-2:])\n",
    "            if segmentations.get('level{:d}'.format(level), None) is None:\n",
    "                segmentations['level{:d}'.format(level)] = []\n",
    "            segmentations['level{:d}'.format(level)].append(label)\n",
    "        segmentations['level{:d}'.format(level)] = torch.stack(segmentations['level{:d}'.format(level)], dim=0)\n",
    "        segmentations['level{:d}'.format(level)] = segmentations['level{:d}'.format(level)].cpu().data.numpy()\n",
    "\n",
    "    # We only visualize the first image in the batch\n",
    "    img = img[0].cpu().data.numpy()\n",
    "    suppixel = suppixel[0].cpu().data.numpy()\n",
    "    img = img * np.array([0.229, 0.224, 0.225])[:, None, None]\n",
    "    img = img + np.array([0.485, 0.456, 0.406])[:, None, None]\n",
    "\n",
    "    # Resize superpixel to the image resolution.\n",
    "    suppixel = cv2.resize(suppixel, (224, 224), interpolation=cv2.INTER_NEAREST)\n",
    "    suppixel = label2color(suppixel, img.transpose(1, 2, 0))\n",
    "\n",
    "    f, ax = plt.subplots(1, 2, figsize=(10, 5));\n",
    "    ax[0].axis('off'); ax[0].imshow(img.transpose(1, 2, 0))\n",
    "    ax[1].axis('off'); ax[1].imshow(suppixel)\n",
    "\n",
    "    f, ax = plt.subplots(1, 4, figsize=(20, 5));\n",
    "    for level in [1, 2, 3, 4]:\n",
    "        seg = segmentations['level{:d}'.format(level)][0]\n",
    "        seg = cv2.resize(seg,\n",
    "                         (224, 224),\n",
    "                         interpolation=cv2.INTER_NEAREST)\n",
    "        ax[level-1].axis('off'); ax[level-1].imshow(CMAP[seg])\n",
    "\n",
    "    # Only iterate once\n",
    "    if i > 0:\n",
    "        break"
   ]
  }
 ],
 "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}