Upload app.py

app.py

@@ -0,0 +1,325 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "78ab80c4-8e25-4464-b710-087d385349fe",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/opt/homebrew/Cellar/jupyterlab/4.4.0/libexec/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n"
+     ]
+    }
+   ],
+   "source": [
+    "import gradio as gr\n",
+    "from PIL import Image\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "import faiss\n",
+    "import json\n",
+    "\n",
+    "from transformers import (\n",
+    "    BlipProcessor,\n",
+    "    BlipForConditionalGeneration,\n",
+    "    CLIPProcessor,\n",
+    "    CLIPModel\n",
+    ")\n",
+    "from datasets import load_dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "9e6fe9c1-df25-41ad-ab27-f6fc20ecb956",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "wikiart_dataset = load_dataset(\"huggan/wikiart\", split=\"train\")\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "b9da3ff0-62e6-4686-af9f-38183f675788",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n"
+     ]
+    }
+   ],
+   "source": [
+    "blip_processor = BlipProcessor.from_pretrained(\"Salesforce/blip-image-captioning-base\")\n",
+    "blip_model = BlipForConditionalGeneration.from_pretrained(\"Salesforce/blip-image-captioning-base\").to(device).eval()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "12d9402a-fdbe-4ade-99ed-26f5d7f9ccfd",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "clip_model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\").to(device).eval()\n",
+    "clip_processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "d4f5e7b2-c873-4495-8ad1-9e32f4f1fbe1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"../create_embeddings/wikiart_embeddings.json\", \"r\", encoding=\"utf-8\") as f:\n",
+    "    data = json.load(f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "87bc4121-f316-4769-bf5d-197db30fe2a3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "image_index = faiss.read_index(\"../create_index/image_index.faiss\")\n",
+    "text_index = faiss.read_index(\"../create_index/text_index.faiss\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "b41d1e5c-d606-4501-a22c-3cde576361d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def generate_caption(image: Image.Image):\n",
+    "    inputs = blip_processor(image, return_tensors=\"pt\").to(device)\n",
+    "    with torch.no_grad():\n",
+    "        caption_ids = blip_model.generate(**inputs)\n",
+    "    caption = blip_processor.decode(caption_ids[0], skip_special_tokens=True)\n",
+    "    return caption"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "263c8672-f4b4-46b7-abf0-483ccfb83c86",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_clip_text_embedding(text):\n",
+    "    inputs = clip_processor(text=[text], return_tensors=\"pt\", padding=True).to(device)\n",
+    "    with torch.no_grad():\n",
+    "        features = clip_model.get_text_features(**inputs)\n",
+    "    features = features.cpu().numpy().astype(\"float32\")\n",
+    "    faiss.normalize_L2(features)\n",
+    "    return features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "34827bd8-e0da-4252-b168-3c79f2d2fb02",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_clip_image_embedding(image):\n",
+    "    inputs = clip_processor(images=image, return_tensors=\"pt\").to(device)\n",
+    "    with torch.no_grad():\n",
+    "        features = clip_model.get_image_features(**inputs)\n",
+    "    features = features.cpu().numpy().astype(\"float32\")\n",
+    "    faiss.normalize_L2(features)\n",
+    "    return features"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "ec6399ac-a40d-49f7-9831-3085fca484c9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def get_results_with_images(embedding, index, top_k=2):\n",
+    "    D, I = index.search(embedding, top_k)\n",
+    "    results = []\n",
+    "    for idx in I[0]:\n",
+    "        item = data[idx]\n",
+    "        img_id = int(item[\"id\"])\n",
+    "        try:\n",
+    "            img = wikiart_dataset[img_id][\"image\"]\n",
+    "        except IndexError:\n",
+    "            continue\n",
+    "        caption = f\"ID: {item['id']}\\n{item['caption']}\"\n",
+    "        results.append((img, caption))\n",
+    "    return results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "76adeb1c-85d6-4e53-9c93-a312c21b71b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def search_similar_images(image: Image.Image):\n",
+    "    caption = generate_caption(image)\n",
+    "\n",
+    "    text_emb = get_clip_text_embedding(caption)\n",
+    "    image_emb = get_clip_image_embedding(image)\n",
+    "\n",
+    "    text_results = get_results_with_images(text_emb, text_index)\n",
+    "    image_results = get_results_with_images(image_emb, image_index)\n",
+    "\n",
+    "    return caption, text_results, image_results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "da86df12-a996-4d1d-ae42-354984cf6dc2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on local URL:  http://127.0.0.1:7862\n",
+      "* To create a public link, set `share=True` in `launch()`.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"http://127.0.0.1:7862/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 13,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "demo = gr.Interface(\n",
+    "    fn=search_similar_images,\n",
+    "    inputs=gr.Image(label=\"Загрузите изображение\", type=\"pil\"),\n",
+    "    outputs=[\n",
+    "        gr.Textbox(label=\"📜 Сгенерированное описание\"),\n",
+    "        gr.Gallery(label=\"🔍 Похожие по описанию (CLIP)\", height=\"auto\", columns=2),\n",
+    "        gr.Gallery(label=\"🎨 Похожие по изображению (CLIP)\", height=\"auto\", columns=2)\n",
+    "    ],\n",
+    "    title=\"🎨 Semantic WikiArt Search (BLIP + CLIP)\",\n",
+    "    description=\"Загрузите изображение. Модель BLIP сгенерирует описание, а CLIP найдёт похожие картины по тексту и изображению.\"\n",
+    ")\n",
+    "\n",
+    "demo.launch()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "55fbac06-4781-4074-a1e6-26ff758bbfe0",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rerunning server... use `close()` to stop if you need to change `launch()` parameters.\n",
+      "----\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
+      "To disable this warning, you can either:\n",
+      "\t- Avoid using `tokenizers` before the fork if possible\n",
+      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "* Running on public URL: https://ba46916423948a3a69.gradio.live\n",
+      "\n",
+      "This share link expires in 1 week. For free permanent hosting and GPU upgrades, run `gradio deploy` from the terminal in the working directory to deploy to Hugging Face Spaces (https://huggingface.co/spaces)\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "<div><iframe src=\"https://ba46916423948a3a69.gradio.live\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": []
+     },
+     "execution_count": 14,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "demo.launch(server_name=\"0.0.0.0\", server_port=7860, share=True)\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c44447c3-0709-4419-a6a4-fc451f80702a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.13.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}