notebooks/eval_ovseg_sam.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3e48b3d-19df-40fd-a85c-d6e675e1f144",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import os\n",
    "import sys\n",
    "from functools import partial\n",
    "\n",
    "import torch\n",
    "import torchvision.transforms as T\n",
    "from torchvision.transforms.functional import InterpolationMode\n",
    "from open_clip import create_model_and_transforms, get_tokenizer\n",
    "\n",
    "import numpy as np\n",
    "from torchmetrics import JaccardIndex\n",
    "from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator\n",
    "\n",
    "from dataset import PartImageNetWithMask, PredictedMask\n",
    "from utils import TextFeatures, get_masked_pred_sam_c, get_masked_pred_sam_f\n",
    "from utils import create_colormap, visualize_img, visualize_seg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d59f14a6-45de-4aed-8dfc-b04bd3718b51",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda:0\"\n",
    "torch.cuda.set_device(device)\n",
    "\n",
    "clip, _, clip_transform = create_model_and_transforms('ViT-B-16', pretrained='openai')\n",
    "tokenizer = get_tokenizer('ViT-B-16')\n",
    "\n",
    "clip = clip.to(device)\n",
    "\n",
    "normalize = clip_transform.transforms[-1]\n",
    "img_transform = T.Compose([\n",
    "    T.Resize(224, interpolation=InterpolationMode.BICUBIC),\n",
    "    T.CenterCrop([224, 224]),\n",
    "])\n",
    "seg_transform = T.Compose([\n",
    "    T.Resize(224, interpolation=InterpolationMode.NEAREST),\n",
    "    T.CenterCrop([224, 224]),\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc38a154-08fc-4793-aab1-d25c92f4d816",
   "metadata": {},
   "outputs": [],
   "source": [
    "# SAM_MODEL = \"vit_b\"\n",
    "# SAM_CKPT_PATH = os.path.join('../sam_ckpt', 'sam_vit_b_01ec64.pth')\n",
    "SAM_MODEL = \"vit_h\"\n",
    "SAM_CKPT_PATH = os.path.join('../sam_ckpt', 'sam_vit_h_4b8939.pth')\n",
    "\n",
    "sam = sam_model_registry[SAM_MODEL](checkpoint=SAM_CKPT_PATH)\n",
    "sam.to(device=device)\n",
    "\n",
    "predictor = SamPredictor(sam)\n",
    "mask_generator = SamAutomaticMaskGenerator(sam)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d74a49a-95b3-4906-9f34-5694bb382a0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_ROOT = '../data/PartImageNet/'\n",
    "SAVE_ROOT = '../pred_segs/'\n",
    "\n",
    "img_root = os.path.join(DATA_ROOT, 'images/val')\n",
    "ano_root = os.path.join(DATA_ROOT, 'annotations/val.json')\n",
    "\n",
    "# Output: image, seg_c, seg_f\n",
    "dataset = PartImageNetWithMask(img_root, ano_root, img_transform, seg_transform)  # use base image transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "656348fd-7eff-4a94-8080-87b4f1328a35",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(dataset.classname_c)\n",
    "print(dataset.classname_f)\n",
    "\n",
    "text_features = TextFeatures(clip, tokenizer, dataset.classname_c, dataset.classname_f)    \n",
    "\n",
    "names = {}\n",
    "for c in dataset.classname_c:\n",
    "    names[c] = [f for f in dataset.classname_f if c in f]\n",
    "cmap_c, cmap_f = create_colormap(names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dc554bb-feb7-485b-b55a-d8de1c3a9359",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_values():\n",
    "    print(\"{:d}/{:d}    {:.2f}/{:.2f}\".format(\n",
    "        index + 1, len(dataset),\n",
    "        np.mean(accs_c) * 100, np.mean(accs_f) * 100,\n",
    "    ))\n",
    "\n",
    "jaccard_c = JaccardIndex(task=\"multiclass\", num_classes=11+1)\n",
    "jaccard_f = JaccardIndex(task=\"multiclass\", num_classes=40+1)\n",
    "\n",
    "accs_c, accs_f = [], []\n",
    "for index in range(len(dataset)):\n",
    "    try:\n",
    "        img_base, seg_c, seg_f = dataset[index]\n",
    "        img = clip_transform(img_base)\n",
    "    \n",
    "        masks = mask_generator.generate(np.array(img_base))\n",
    "        masks = sorted(masks, key=lambda m: m[\"area\"], reverse=True)\n",
    "        masks = [torch.from_numpy(m['segmentation']).unsqueeze(0) for m in masks]\n",
    "        mask_c = mask_f = masks\n",
    "    \n",
    "        pred_c = get_masked_pred_sam_c(clip, text_features, img, mask_c)\n",
    "        pred_f = get_masked_pred_sam_f(clip, text_features, img, mask_f, pred_c)\n",
    "\n",
    "        accs_c.append(jaccard_c(pred_c, seg_c).item())\n",
    "        accs_f.append(jaccard_f(pred_f, seg_f).item())\n",
    "    except:\n",
    "        pass\n",
    "    \n",
    "    if (index + 1) % 100 == 0:\n",
    "        print_values()\n",
    "\n",
    "index = len(dataset) - 1\n",
    "print_values()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9ab797c-9226-47e6-a9be-b64a892a6e1b",
   "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
}