clip.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "colab": {
      "name": "clip.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "##Semantic image search with OpenAI's CLIP\n",
        "Demo of [OpenAI's CLIP](https://openai.com/blog/clip/):\n",
        "- built with [transformers](https://huggingface.co/transformers/) from \n",
        "🤗 Hugging Face\n",
        "- based on 25,000 images from [Unsplash](https://unsplash.com/) and 7,685 images from the [Movie Database (TMDB)](https://www.themoviedb.org/)\n",
        "- inspired by [Unsplash Image Search](https://github.com/haltakov/natural-language-image-search) by Vladimir Haltakov and [Alph, The Sacred River](https://github.com/thoppe/alph-the-sacred-river) by Travis Hoppe "
      ],
      "metadata": {
        "id": "Gk-ai6za6Apu"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [
        "#@title ##← Click on the circled arrow and wait for up to 1 minute\n",
        "!pip install -q transformers >/dev/null\n",
        "\n",
        "import pandas as pd, numpy as np\n",
        "import os\n",
        "import urllib.request\n",
        "from transformers import CLIPProcessor, CLIPTextModel, CLIPModel, logging\n",
        "\n",
        "from IPython.display import display, Markdown, HTML, clear_output\n",
        "import ipywidgets as widgets\n",
        "\n",
        "logging.get_verbosity = lambda: logging.NOTSET\n",
        "model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "clear_output()\n",
        "\n",
        "def compute_text_embeddings(list_of_strings):\n",
        "    inputs = processor(text=list_of_strings, return_tensors=\"pt\", padding=True)\n",
        "    return model.get_text_features(**inputs)\n",
        "\n",
        "urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1bt1O-iArKuU9LGkMV1zUPTEHZk8k7L65', 'data.csv')\n",
        "urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=19aVnFBY-Rc0-3VErF_C7PojmWpBsb5wk', 'data2.csv')\n",
        "urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1onKr-pfWb4l6LgL-z8WDod3NMW-nIJxE', 'embeddings.npy')\n",
        "urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1KbwUkE0T8bpnHraqSzTeGGV4-TZO_CFB', 'embeddings2.npy')\n",
        "\n",
        "df = {0: pd.read_csv('data.csv'), 1: pd.read_csv('data2.csv')}\n",
        "embeddings = {0: np.load('embeddings.npy'), 1: np.load('embeddings2.npy')}\n",
        "for k in [0, 1]:\n",
        "  embeddings[k] = np.divide(embeddings[k], np.sqrt(np.sum(embeddings[k]**2, axis=1, keepdims=True)))\n",
        "source = {0: '\\nSource: Unsplash', 1: '\\nSource: The Movie Database (TMDB)'}\n",
        "\n",
        "def get_html(url_list, height=200):\n",
        "    html = \"<div style='margin-top: 20px; display: flex; flex-wrap: wrap; justify-content: space-evenly'>\"\n",
        "    for url, title, link in url_list:\n",
        "        html2 = f\"<img title='{title}' style='height: {height}px; margin-bottom: 10px' src='{url}'>\"\n",
        "        if len(link) > 0:\n",
        "            html2 = f\"<a href='{link}' target='_blank'>\" + html2 + \"</a>\"\n",
        "        html = html + html2\n",
        "    html += \"</div>\"\n",
        "    return html\n",
        "\n",
        "query = widgets.Text(layout=widgets.Layout(width='400px'))\n",
        "dataset =widgets.Dropdown(\n",
        "    options=['Unsplash', 'Movies'],\n",
        "    value='Unsplash'\n",
        ")\n",
        "button = widgets.Button(description=\"Search\")\n",
        "output = widgets.Output()\n",
        "\n",
        "display(widgets.HBox([query, button, dataset],\n",
        "                     layout=widgets.Layout(justify_content='center')),\n",
        "        output)\n",
        "\n",
        "def image_search(query, n_results=24):\n",
        "    text_embeddings = compute_text_embeddings([query]).detach().numpy()\n",
        "    k = 0 if dataset.value == 'Unsplash' else 1\n",
        "    results = np.argsort((embeddings[k]@text_embeddings.T)[:, 0])[-1:-n_results-1:-1]\n",
        "    return [(df[k].iloc[i]['path'],\n",
        "             df[k].iloc[i]['tooltip'] + source[k],\n",
        "             df[k].iloc[i]['link']) for i in results]\n",
        "\n",
        "def on_button_clicked(b):\n",
        "    if len(query.value) > 0:\n",
        "        results = image_search(query.value)\n",
        "        output.clear_output()\n",
        "        with output:\n",
        "            display(HTML(get_html(results)))\n",
        "\n",
        "button.on_click(on_button_clicked)\n",
        "dataset.observe(on_button_clicked, names='value')"
      ],
      "outputs": [],
      "metadata": {
        "id": "uzdFhRGqiWkR",
        "cellView": "form"
      }
    }
  ]
}