{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8cd1e865-53d5-460b-8bae-5658e3aa3d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "import panel as pn\n",
    "pn.extension()\n",
    "import requests\n",
    "import random\n",
    "import PIL\n",
    "from PIL import Image\n",
    "import io\n",
    "from transformers import CLIPProcessor, CLIPModel\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8570053-0b83-421b-95c2-695b6c709ba1",
   "metadata": {},
   "outputs": [],
   "source": [
    "pn.extension('texteditor', template=\"bootstrap\", sizing_mode='stretch_width')\n",
    "\n",
    "pn.state.template.param.update(\n",
    "    main_max_width=\"690px\",\n",
    "    header_background=\"#F08080\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca65cc07-8181-4259-8770-9c780621eb78",
   "metadata": {},
   "outputs": [],
   "source": [
    "# File input widget\n",
    "file_input = pn.widgets.FileInput()\n",
    "\n",
    "# Button widget\n",
    "compute_button = pn.widgets.Button(name=\"Compute\")\n",
    "\n",
    "# Text input widget\n",
    "text_input = pn.widgets.TextInput(name='Possible class names (e.g., cat, dog)', placeholder='cat, dog')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3691594-df8c-4d03-99e8-db4d3b2520c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_image(value, width=600):\n",
    "    \"\"\"\n",
    "    normalize image to RBG channels and to the same size\n",
    "    \"\"\"\n",
    "    if value: \n",
    "        b = io.BytesIO(value)\n",
    "        image = PIL.Image.open(b).convert(\"RGB\")\n",
    "    else: \n",
    "        url = \"http://images.cocodataset.org/val2017/000000039769.jpg\"\n",
    "        image = Image.open(requests.get(url, stream=True).raw)\n",
    "    aspect = image.size[1] / image.size[0]\n",
    "    height = int(aspect * width)\n",
    "    return image.resize((width, height), PIL.Image.LANCZOS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b139802-c9d6-4493-acb2-5051343c1ecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def image_classification(image):\n",
    "    model = CLIPModel.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "    processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-large-patch14\")\n",
    "    possible_categories = text_input.value.split(\",\")\n",
    "    if text_input.value == '':\n",
    "        possible_categories = ['cat', ' dog']\n",
    "    inputs = processor(text=possible_categories, images=image, return_tensors=\"pt\", padding=True)\n",
    "    \n",
    "    outputs = model(**inputs)\n",
    "    logits_per_image = outputs.logits_per_image # this is the image-text similarity score\n",
    "    probs = logits_per_image.softmax(dim=1)\n",
    "    return probs.detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b6f0ce5-03a5-4a14-b0b7-74c8190ce928",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_result(_):\n",
    "    image = normalize_image(file_input.value)\n",
    "\n",
    "    result = image_classification(image)\n",
    "    \n",
    "    possible_categories = text_input.value.split(\",\")\n",
    "    if text_input.value == '':\n",
    "        possible_categories = ['cat', ' dog']\n",
    "\n",
    "    progress_bars = pn.Column(*[\n",
    "        pn.Row(\n",
    "            possible_categories[i], \n",
    "            pn.indicators.Progress(name='', value=int(j*100), width=500))\n",
    "        for i, j in enumerate(result[0])\n",
    "    ])\n",
    "    return progress_bars\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fd5a63f-012a-419c-8386-22b5b8ff243f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Bind the get_image function with the button widget\n",
    "interactive_result = pn.bind(get_result, compute_button)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "399189f1-4ff6-4f4b-b050-76e9a46443dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# layout\n",
    "pn.Column(\n",
    "    \"## \\U0001F60A Upload an image file and start classifying!\",\n",
    "    file_input,\n",
    "    pn.bind(pn.panel, file_input),\n",
    "    text_input, \n",
    "    compute_button,\n",
    "    interactive_result\n",
    ").servable(title=\"Panel Image Classification Demo\")"
   ]
  }
 ],
 "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.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}