notebooks/save_segmentation_cast.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": "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": [
    "DATA_ROOT = '../data/PartImageNet/images/val/'\n",
    "# DATA_ROOT = './demo_images'\n",
    "SAVE_ROOT = '../pred_segs/'\n",
    "\n",
    "class ReturnIndexDataset(datasets.ImageFolder):\n",
    "    def __getitem__(self, idx):\n",
    "        img, seg, lab = super(ReturnIndexDataset, self).__getitem__(idx)\n",
    "        return img, seg, lab, idx\n",
    "\n",
    "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 = ReturnIndexDataset(\n",
    "    DATA_ROOT,\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, img_inds) 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).detach()\n",
    "            if level == 1 and len(label) < 196:\n",
    "                npad = torch.unique(suppixel).shape[0] - logit.shape[0]\n",
    "                label = torch.concat([label, torch.tensor([label[-1]] * npad, device=label.device)])\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",
    "\n",
    "            # Save segmentations by levels\n",
    "            save_root = os.path.join(SAVE_ROOT, MODEL_CLASS_NAME, 'level{:d}'.format(level))\n",
    "            os.makedirs(save_root, exist_ok=True)\n",
    "\n",
    "            file_name = train_dataset.samples[img_inds[b]][0]\n",
    "            file_name = file_name.split('/')[-1].split('.')[0]\n",
    "            with open(os.path.join(save_root, '{}.npy'.format(file_name)), 'wb') as f:\n",
    "                np.save(f, label.cpu().data.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0ab1930-42c0-43a3-b9fa-0b156e8d64f5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}