{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bb9b2af6-325f-4d1e-8e74-96ca5c2e27c5",
   "metadata": {},
   "source": [
    "# OpenClip\n",
    "\n",
    "[OpenClip](https://github.com/mlfoundations/open_clip/tree/main) is an source implementation of OpenAI's CLIP.\n",
    "\n",
    "These multi-modal embeddings can be used to embed images or text."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64ce4cbd-d81b-4278-a33b-a38e6a9a5de4",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U langchain-experimental"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "570d818f-5705-4532-8e77-b0f335bb515d",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install pillow open_clip_torch torch matplotlib"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d5307f2-37c0-49fc-9337-e5296189c240",
   "metadata": {},
   "source": [
    "We can the list of available CLIP embedding models and checkpoints:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50de46c3-6e6d-4fab-b6a0-a42e2a7a458e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import open_clip\n",
    "\n",
    "open_clip.list_pretrained()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f25dbe1-0ebc-4b23-8b77-301e9c811fe3",
   "metadata": {},
   "source": [
    "Below, I test a larger but more performant model based on the table ([here](https://github.com/mlfoundations/open_clip)):\n",
    "```\n",
    "model_name = \"ViT-g-14\"\n",
    "checkpoint = \"laion2b_s34b_b88k\"\n",
    "```\n",
    "\n",
    "But, you can also opt for a smaller, less performant model:\n",
    "```\n",
    "model_name = \"ViT-B-32\"\n",
    "checkpoint = \"laion2b_s34b_b79k\"\n",
    "```\n",
    "\n",
    "The model `model_name`,`checkpoint`  are set in `langchain_experimental.open_clip.py`.\n",
    "\n",
    "For text, use the same method `embed_documents` as with other embedding models.\n",
    "\n",
    "For images, use `embed_image` and simply pass a list of uris for the images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c4fc9af7-5659-4008-b6a3-0f99c68324aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from langchain_experimental.open_clip import OpenCLIPEmbeddings\n",
    "from PIL import Image\n",
    "\n",
    "# Image URIs\n",
    "uri_dog = \"/Users/rlm/Desktop/test/dog.jpg\"\n",
    "uri_house = \"/Users/rlm/Desktop/test/house.jpg\"\n",
    "\n",
    "# Embe images or text\n",
    "clip_embd = OpenCLIPEmbeddings()\n",
    "img_feat_dog = clip_embd.embed_image([uri_dog])\n",
    "img_feat_house = clip_embd.embed_image([uri_house])\n",
    "text_feat_dog = clip_embd.embed_documents([\"dog\"])\n",
    "text_feat_house = clip_embd.embed_documents([\"house\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f732e787-e12f-4047-a0b3-2062a2471cff",
   "metadata": {},
   "source": [
    "## Sanity Check\n",
    "\n",
    "Let's reproduce results shown in the OpenClip Colab [here](https://colab.research.google.com/github/mlfoundations/open_clip/blob/master/docs/Interacting_with_open_clip.ipynb#scrollTo=tMc1AXzBlhzm)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "02a14a40-d0a5-44ce-b14c-1005f8da959b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 2930188 characters]",
      "text/plain": [
       "<Figure size 1600x500 with 8 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 487,
       "width": 1539
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import os\n",
    "from collections import OrderedDict\n",
    "\n",
    "import IPython.display\n",
    "import matplotlib.pyplot as plt\n",
    "import skimage\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina'\n",
    "\n",
    "descriptions = {\n",
    "    \"page\": \"a page of text about segmentation\",\n",
    "    \"chelsea\": \"a facial photo of a tabby cat\",\n",
    "    \"astronaut\": \"a portrait of an astronaut with the American flag\",\n",
    "    \"rocket\": \"a rocket standing on a launchpad\",\n",
    "    \"motorcycle_right\": \"a red motorcycle standing in a garage\",\n",
    "    \"camera\": \"a person looking at a camera on a tripod\",\n",
    "    \"horse\": \"a black-and-white silhouette of a horse\",\n",
    "    \"coffee\": \"a cup of coffee on a saucer\",\n",
    "}\n",
    "\n",
    "original_images = []\n",
    "images = []\n",
    "image_uris = []  # List to store image URIs\n",
    "texts = []\n",
    "plt.figure(figsize=(16, 5))\n",
    "\n",
    "# Loop to display and prepare images and assemble URIs\n",
    "for filename in [\n",
    "    filename\n",
    "    for filename in os.listdir(skimage.data_dir)\n",
    "    if filename.endswith(\".png\") or filename.endswith(\".jpg\")\n",
    "]:\n",
    "    name = os.path.splitext(filename)[0]\n",
    "    if name not in descriptions:\n",
    "        continue\n",
    "\n",
    "    image_path = os.path.join(skimage.data_dir, filename)\n",
    "    image = Image.open(image_path).convert(\"RGB\")\n",
    "\n",
    "    plt.subplot(2, 4, len(images) + 1)\n",
    "    plt.imshow(image)\n",
    "    plt.title(f\"{filename}\\n{descriptions[name]}\")\n",
    "    plt.xticks([])\n",
    "    plt.yticks([])\n",
    "\n",
    "    original_images.append(image)\n",
    "    images.append(image)  # Origional code does preprocessing here\n",
    "    texts.append(descriptions[name])\n",
    "    image_uris.append(image_path)  # Add the image URI to the list\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1294325d-1525-40ce-b445-f4bf89da1611",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, 'Cosine similarity between text and image features')"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 1068340 characters]",
      "text/plain": [
       "<Figure size 2000x1400 with 1 Axes>"
      ]
     },
     "metadata": {
      "image/png": {
       "height": 1127,
       "width": 1459
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Instantiate your model\n",
    "clip_embd = OpenCLIPEmbeddings()\n",
    "\n",
    "# Embed images and text\n",
    "img_features = clip_embd.embed_image(image_uris)\n",
    "text_features = clip_embd.embed_documents([\"This is \" + desc for desc in texts])\n",
    "\n",
    "# Convert the list of lists to numpy arrays for matrix operations\n",
    "img_features_np = np.array(img_features)\n",
    "text_features_np = np.array(text_features)\n",
    "\n",
    "# Calculate similarity\n",
    "similarity = np.matmul(text_features_np, img_features_np.T)\n",
    "\n",
    "# Plot\n",
    "count = len(descriptions)\n",
    "plt.figure(figsize=(20, 14))\n",
    "plt.imshow(similarity, vmin=0.1, vmax=0.3)\n",
    "# plt.colorbar()\n",
    "plt.yticks(range(count), texts, fontsize=18)\n",
    "plt.xticks([])\n",
    "for i, image in enumerate(original_images):\n",
    "    plt.imshow(image, extent=(i - 0.5, i + 0.5, -1.6, -0.6), origin=\"lower\")\n",
    "for x in range(similarity.shape[1]):\n",
    "    for y in range(similarity.shape[0]):\n",
    "        plt.text(x, y, f\"{similarity[y, x]:.2f}\", ha=\"center\", va=\"center\", size=12)\n",
    "\n",
    "for side in [\"left\", \"top\", \"right\", \"bottom\"]:\n",
    "    plt.gca().spines[side].set_visible(False)\n",
    "\n",
    "plt.xlim([-0.5, count - 0.5])\n",
    "plt.ylim([count + 0.5, -2])\n",
    "\n",
    "plt.title(\"Cosine similarity between text and image features\", size=20)"
   ]
  }
 ],
 "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.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
