CLIP_Patch_Detection.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/kevinzakka/clip_playground/blob/main/CLIP_Patch_Detection.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QpPbEKLBmyX1"
   },
   "source": [
    "# CLIP Patch Detection\n",
    "\n",
    "This Colab notebook demos crude object detection by spliting an image into patches and finding the highest patch-caption similarity in CLIP embedding space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "qQOvOhnKQ-Tu"
   },
   "outputs": [],
   "source": [
    "#@title Install dependencies\n",
    "\n",
    "#@markdown Please execute this cell by pressing the _Play_ button \n",
    "#@markdown on the left.\n",
    "\n",
    "#@markdown **Note**: This installs the software on the Colab \n",
    "#@markdown notebook in the cloud and not on your computer.\n",
    "\n",
    "%%capture\n",
    "!pip install ftfy regex tqdm matplotlib\n",
    "!pip install git+https://github.com/openai/CLIP.git\n",
    "\n",
    "import math\n",
    "import urllib.request\n",
    "import clip\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from PIL import Image\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "caPbAhFlRBwT"
   },
   "outputs": [],
   "source": [
    "#@title Helper functions\n",
    "\n",
    "#@markdown Some helper functions for loading, patchifying and visualizing images.\n",
    "\n",
    "def load_image(img_path, resize=None, pil=False):\n",
    "    image = Image.open(img_path).convert(\"RGB\")\n",
    "    if resize is not None:\n",
    "        image = image.resize((resize, resize))\n",
    "    if pil:\n",
    "        return image\n",
    "    return np.asarray(image).astype(np.float32) / 255.\n",
    "\n",
    "def viz_patches(x, figsize=(20, 20), patch_idx=None, topk=None, t=5):\n",
    "    # x: num_patches, 3, patch_size, patch_size\n",
    "    n = x.shape[0]\n",
    "    nrows = int(math.sqrt(n))\n",
    "    _, axes = plt.subplots(nrows, nrows, figsize=figsize)\n",
    "    for i, ax in enumerate(axes.flatten()):            \n",
    "        im = x[i].permute(1, 2, 0).numpy()\n",
    "        im = (im * 255.).round().astype(np.uint8)\n",
    "        if patch_idx is not None and i == patch_idx:\n",
    "            im[0:t] = (255, 0, 0)\n",
    "            im[im.shape[0]-t:] = (255, 0, 0)\n",
    "            im[:, 0:t] = (255, 0, 0)\n",
    "            im[:, im.shape[1]-t:] = (255, 0, 0)\n",
    "        if topk is not None:\n",
    "            if i in topk and i != patch_idx:\n",
    "                im[0:t] = (255, 255, 0)\n",
    "                im[im.shape[0]-t:] = (255, 255, 0)\n",
    "                im[:, 0:t] = (255, 255, 0)\n",
    "                im[:, im.shape[1]-t:] = (255, 255, 0)\n",
    "        ax.imshow(im)\n",
    "        ax.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "def patchify(image_path, resolution, patch_size, patch_stride=None):\n",
    "    img_tensor = transforms.ToTensor()(load_image(image_path, resolution, True))\n",
    "    if patch_stride is None:\n",
    "        patch_stride = patch_size\n",
    "    patches = img_tensor.unfold(\n",
    "        1, patch_size, patch_stride).unfold(2, patch_size, patch_stride)\n",
    "    patches = patches.reshape(3, -1, patch_size, patch_size).permute(1, 0, 2, 3)\n",
    "    return patches  # N, 3, patch_size, patch_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 482
    },
    "id": "hwmmPfnuemkD",
    "outputId": "96c39215-b829-4d75-9d0d-087cd3c26085"
   },
   "outputs": [],
   "source": [
    "#@title Image and Patch Settings { run: \"auto\" }\n",
    "\n",
    "image_url = 'https://images2.minutemediacdn.com/image/upload/c_crop,h_706,w_1256,x_0,y_64/f_auto,q_auto,w_1100/v1554995050/shape/mentalfloss/516438-istock-637689912.jpg' #@param {type:\"string\"}\n",
    "image_resolution =  900#@param {type:\"integer\"}\n",
    "patch_size =  224#@param {type:\"integer\"}\n",
    "# integer_input = 10 #@param {type:\"integer\"}\n",
    "# integer_slider = 21 #@param {type:\"slider\", min:0, max:100, step:1}\n",
    "\n",
    "# Download the image from the web.\n",
    "image_path = 'image.png'\n",
    "urllib.request.urlretrieve(image_url, image_path)\n",
    "\n",
    "patches = patchify(image_path, image_resolution, patch_size)\n",
    "print(\"patches: \", patches.shape)\n",
    "viz_patches(patches, figsize=(8, 8))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 575
    },
    "id": "QuHrrEqXdVmx",
    "outputId": "c0adb805-22a9-4353-d98e-65bd1b3e8772"
   },
   "outputs": [],
   "source": [
    "#@title Detect\n",
    "\n",
    "clip_model = \"RN50\" #@param [\"RN50\", \"RN101\", \"RN50x4\", \"RN50x16\", \"ViT-B/32\", \"ViT-B/16\"]\n",
    "image_caption = 'the dog' #@param {type:\"string\"}\n",
    "topk =  6#@param {type:\"integer\"}\n",
    "\n",
    "# Load CLIP model.\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model, preprocess = clip.load(clip_model, device=device, jit=False)\n",
    "\n",
    "text_input = clip.tokenize([image_caption]).to(device)\n",
    "\n",
    "# Pad in case not equal to model expected input resolution.\n",
    "p = model.visual.input_resolution - patch_size\n",
    "patches_pad = torch.nn.functional.pad(\n",
    "    patches, (p//2, p//2, p//2, p//2), \"constant\", 0).to(device)\n",
    "\n",
    "with torch.no_grad():\n",
    "    patch_embs = model.encode_image(patches_pad)\n",
    "    text_embs = model.encode_text(text_input)\n",
    "    patch_embs = patch_embs / patch_embs.norm(dim=-1, keepdim=True)\n",
    "    text_embs = text_embs / text_embs.norm(dim=-1, keepdim=True)\n",
    "    sim = patch_embs @ text_embs.t()\n",
    "    idx_max = sim.argmax().item()\n",
    "    topk_idxs = torch.topk(sim.flatten(), topk)[-1].cpu().numpy().tolist()\n",
    "\n",
    "viz_patches(patches, figsize=(10, 10), patch_idx=idx_max, topk=topk_idxs, t=int(0.05*patch_size))"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "CLIP Patch Detection",
   "provenance": []
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}