DoomAI/doomai

1import gradio as gr2from PIL import Image3from datasets import load_dataset4from transformers import CLIPProcessor, CLIPModel5from annoy import AnnoyIndex6import json7import torch8import cv29 10model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")11processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")12 13# Загружаем Annoy-индекс14embedding_dim = 51215index = AnnoyIndex(embedding_dim, 'angular')16index.load("image_embeddings_256.ann")17 18# Загружаем метаданные19with open("annoy_metadata.json", 'r') as f:20    metadata = json.load(f)21 22 23def find_box(input_image):24    image = cv2.imread(input_image)25 26    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)27    gray = cv2.blur(gray, (7, 7), 0)28 29    # Применим метод Canny для нахождения контуров30    edges = cv2.Canny(gray, 70, 130)31 32    # Увеличим контуры для лучшей детекции33    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))34    dilated = cv2.dilate(edges, kernel, iterations=2)35 36    # Найдем контуры на изображении37    contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)38 39    # Нарисуем контуры40    output = image.copy()41 42    h_glob, w_glob, _ = image.shape43    s_glob = h_glob * w_glob44    max_s = 045    found = None46 47    for contour in contours:48        # Получим прямоугольник, обрамляющий контур49        x, y, w, h = cv2.boundingRect(contour)50        s = h * w51        # Фильтруем по размеру, чтобы убрать мелкие шумы52        if w > 50 and h > 50 and s <= 0.7 * s_glob and 0.3 < w / h < 3 and s > max_s:53            cv2.rectangle(output, (x, y), (x + w, y + h), (0, 255, 0), 2)54            max_s = s55            found = image[y:y + h, x:x + w]56 57    if found is None:58        return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)59    return cv2.cvtColor(found, cv2.COLOR_BGR2RGB)60 61 62def run(text, input_image, n_pictures, detect_flag):63    if not input_image:64        return65 66    if detect_flag:67        image = Image.fromarray(find_box(input_image)).convert("RGB")68    else:69        image = Image.open(input_image).convert("RGB")70 71    # Преобразуем изображение в вектор с помощью CLIP72    inputs = processor(images=image, return_tensors="pt")73    with torch.no_grad():74        image_embedding = model.get_image_features(**inputs)75        image_embedding = image_embedding.cpu().numpy().flatten()76 77    final_embedding = image_embedding78 79    text = text.strip()80    if text:81        inputs = processor(text=text, return_tensors="pt")82        with torch.no_grad():83            text_embedding = model.get_text_features(**inputs)84            text_embedding = text_embedding.cpu().numpy().flatten()85 86        final_embedding += text_embedding * 0.487 88    indices = index.get_nns_by_vector(final_embedding, n_pictures)89 90    files = []91    for idx in indices:92        # Индексы хранятся как строки, поэтому преобразуем их в строку93        metadata_entry = metadata.get(str(idx))94        files.append(f"{metadata_entry['directory']}/{metadata_entry['filename']}")95 96    return load_dataset("DoomAI/doomai", data_files=files)["train"]["image"]97 98 99# Gradio100with gr.Blocks(101        theme=gr.themes.Default(primary_hue=gr.themes.colors.blue, secondary_hue=gr.themes.colors.orange)) as demo:102    gr.Markdown("# DoomAI")103    with gr.Row():104        with gr.Column():105            # Ввод изображения106            image_input = gr.Image(label="Картинка", type="filepath", sources=['upload', 'clipboard'])107 108            # Дополнительные параметры109            pictures_slider = gr.Slider(110                label="Количество картинок", minimum=1.0, maximum=20.0, value=10, step=1,111            )112            checkbox = gr.Checkbox(113                False, label="Удаление фона скриншота"114            )115            prompt_input = gr.Textbox(116                label="Уточнение изображения на английском (опционально)", placeholder="A photo of a cute dog"117            )118 119            # Кнопка120            generate_button = gr.Button("Найти", variant='primary')121 122        with gr.Column():123            # Вывод124            gallery = gr.Gallery(label='Галерея', columns=4, rows=5, show_share_button=False)125 126    # При нажатии127    generate_button.click(128        run,129        inputs=[130            prompt_input,131            image_input,132            pictures_slider,133            checkbox134        ],135        outputs=gallery136    )137 138# Запуск приложения139demo.launch()140