| { | |
| "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 | |
| } | |