Upload app.py

app.py

@@ -1,325 +1,95 @@
-{
- "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
-}
+import gradio as gr
+from PIL import Image
+import torch
+import numpy as np
+import faiss
+
+from transformers import (
+    BlipProcessor,
+    BlipForConditionalGeneration,
+    CLIPProcessor,
+    CLIPModel
+)
+from datasets import load_dataset
+
+# Загрузка датасета и моделей
+wikiart_dataset = load_dataset("huggan/wikiart", split="train")
+device = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu")
+
+blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(device).eval()
+
+clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device).eval()
+clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+
+# Загрузка FAISS индексов
+image_index = faiss.read_index("image_index.faiss")
+text_index = faiss.read_index("text_index.faiss")
+
+# Генерация описания через BLIP
+def generate_caption(image: Image.Image):
+    inputs = blip_processor(image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        caption_ids = blip_model.generate(**inputs)
+    caption = blip_processor.decode(caption_ids[0], skip_special_tokens=True)
+    return caption
+
+# Получение CLIP эмбеддинга по тексту
+def get_clip_text_embedding(text):
+    inputs = clip_processor(text=[text], return_tensors="pt", padding=True).to(device)
+    with torch.no_grad():
+        features = clip_model.get_text_features(**inputs)
+    features = features.cpu().numpy().astype("float32")
+    faiss.normalize_L2(features)
+    return features
+
+# Получение CLIP эмбеддинга по изображению
+def get_clip_image_embedding(image):
+    inputs = clip_processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        features = clip_model.get_image_features(**inputs)
+    features = features.cpu().numpy().astype("float32")
+    faiss.normalize_L2(features)
+    return features
+
+# Получение похожих изображений по эмбеддингу
+def get_results_with_images(embedding, index, top_k=2):
+    D, I = index.search(embedding, top_k)
+    results = []
+    for idx in I[0]:
+        try:
+            item = wikiart_dataset[idx]
+            img = item["image"]
+            title = item.get("title", "Untitled")
+            artist = item.get("artist", "Unknown")
+            caption = f"ID: {idx}\n{title} — {artist}"
+            results.append((img, caption))
+        except IndexError:
+            continue
+    return results
+
+# Основная функция поиска
+def search_similar_images(image: Image.Image):
+    caption = generate_caption(image)
+    text_emb = get_clip_text_embedding(caption)
+    image_emb = get_clip_image_embedding(image)
+
+    text_results = get_results_with_images(text_emb, text_index)
+    image_results = get_results_with_images(image_emb, image_index)
+
+    return caption, text_results, image_results
+
+# Интерфейс Gradio
+demo = gr.Interface(
+    fn=search_similar_images,
+    inputs=gr.Image(label="Загрузите изображение", type="pil"),
+    outputs=[
+        gr.Textbox(label="📜 Сгенерированное описание"),
+        gr.Gallery(label="🔍 Похожие по описанию (CLIP)", height="auto", columns=2),
+        gr.Gallery(label="🎨 Похожие по изображению (CLIP)", height="auto", columns=2)
+    ],
+    title="🎨 Semantic WikiArt Search (BLIP + CLIP)",
+    description="Загрузите изображение. Модель BLIP сгенерирует описание, а CLIP найдёт похожие картины по тексту и изображению."
+)
+
+demo.launch()