Zero_Shot_Segmentation_Clip.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3dc78803",
   "metadata": {},
   "source": [
    "## Zero Shot Segmentation using OPEN AI's CLIP"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2e2fdbff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torchvision import transforms\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import clip\n",
    "from PIL import Image\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c15089d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model, preprocess = clip.load(\"ViT-B/32\", device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2d23e4cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from encoder import Encoder\n",
    "from decoder import Decoder\n",
    "from dataloader import PhraseCutDataset_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "327b8c92",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "encoder = Encoder()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d9e8b27",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "print(encoder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef45f940",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = PhraseCutDataset_('val')\n",
    "val = DataLoader(data, batch_size=1, shuffle=False)\n",
    "\n",
    "data = PhraseCutDataset_('train')\n",
    "train = DataLoader(data, batch_size=1, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beaf7b0d",
   "metadata": {},
   "source": [
    "### Decoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "65757411",
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder = Decoder(extract_layers=[3, 6, 9], mha_heads=4, reduce_dim=64, cond_layer=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17061376",
   "metadata": {},
   "outputs": [],
   "source": [
    "decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99f84e68",
   "metadata": {},
   "source": [
    "### DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3a1c90fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ca0f8a17",
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.BCEWithLogitsLoss()\n",
    "optimizer = torch.optim.AdamW(decoder.parameters(), lr=3e-4)\n",
    "lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 307480, 1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "793c0206",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88017dd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_hist = []\n",
    "iter_n = []\n",
    "epochs = 10\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    \n",
    "    for i, (phrase, input_img, output_img, id_) in enumerate(train):\n",
    "    \n",
    "        accuracies_iou = 0\n",
    "\n",
    "        if(len(input_img.shape) != 4):\n",
    "            continue\n",
    "\n",
    "        encodings = encoder(transforms.ToPILImage()(input_img[0].permute(2, 0, 1)).convert(\"RGB\"), phrase[0])\n",
    "\n",
    "        output = decoder(encodings)\n",
    "\n",
    "        loss = criterion(output[0][0], output_img[0])\n",
    "\n",
    "        pred = (torch.sigmoid(output[0][0]))#>0.3).int()\n",
    "\n",
    "        tp = torch.sum(pred*output_img[0])\n",
    "        fp = torch.sum(pred*(1. - output_img[0]))\n",
    "        fn = torch.sum((1. - pred)*output_img[0])\n",
    "        accuracies_iou *= i\n",
    "        accuracies_iou += (tp/(tp+fp+fn))\n",
    "        accuracies_iou /= i + 1\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "\n",
    "        if (i+1)%(3000) == 0:\n",
    "            loss_hist.append(loss.item())\n",
    "            iter_n.append(epoch*30748 + i)\n",
    "            print(f\"Epoch : {epoch + 1}, Iteration : {i+1}, Loss : {loss.item()}\")\n",
    "            print(f\"IOU Accuracy : {100*(accuracies_iou) :.5f}%\")\n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "4ea92853",
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(decoder.state_dict(), \"ClipSeg.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee6de5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(iter_n, loss_hist)\n",
    "plt.title(\"Loss vs Iterations\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49340ecb",
   "metadata": {},
   "source": [
    "### Model Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52f63744",
   "metadata": {},
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    pixel = []\n",
    "    iou = [] \n",
    "    \n",
    "    for i, (phrase, input_img, output_img, _) in enumerate(val):\n",
    "        \n",
    "        if(len(input_img.shape) != 4):\n",
    "            continue\n",
    "\n",
    "        encodings = encoder(transforms.ToPILImage()(input_img[0].permute(2, 0, 1)).convert(\"RGB\"), phrase[0])\n",
    "\n",
    "        output = decoder(encodings)\n",
    "        \n",
    "        pred = ((torch.sigmoid(output[0][0])) > 0.3).float()\n",
    "        \n",
    "        pixel.append(torch.sum((pred) == output_img[0])/(224*224))\n",
    "        \n",
    "        tp = torch.sum(pred*output_img[0])\n",
    "        fp = torch.sum(pred*(1. - output_img[0]))\n",
    "        fn = torch.sum((1. - pred)*output_img[0])\n",
    "\n",
    "        iou.append(tp/(tp+fp+fn))\n",
    "        \n",
    "    print(f\"Pixel-by-Pixel Accuracy : {100*sum(pixel)/len(pixel) :.5f}%\")\n",
    "    print(f\"IOU Accuracy : {100*sum(iou)/len(iou) :.5f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "692a83b2",
   "metadata": {},
   "source": [
    "# Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61c831a",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Decoder(extract_layers=[3, 6, 9], mha_heads=4, cond_layer=3, reduce_dim=64)\n",
    "model.load_state_dict(torch.load('__pycache__/ClipSeg.pth'))\n",
    "for param in model.parameters():\n",
    "    param.requires_grad_(False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a17fcd7",
   "metadata": {},
   "outputs": [],
   "source": [
    "img=Image.open(\"PhraseCutDataset/data/VGPhraseCut_v0/images_val/2339423.jpg\")\n",
    "\n",
    "img = img.resize((224,224))\n",
    "img.show()\n",
    "\n",
    "encodings = encoder(img, \"stack of gifts\")\n",
    "output = model(encodings)\n",
    "\n",
    "pred=(torch.sigmoid(output)>0.3).float()\n",
    "\n",
    "img = transforms.ToPILImage()(pred[0]).convert(\"L\")\n",
    "img.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}