import gradio as gr
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import torch

# Загружаем модель и препроцессор
model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")

# Основная функция для классификации
def classify(image, labels):
    # labels — строка вида "Titanic, Avatar, Interstellar"
    candidate_labels = [label.strip() for label in labels.split(",") if label.strip()]
    inputs = processor(text=candidate_labels, images=image, return_tensors="pt", padding=True)

    with torch.no_grad():
        outputs = model(**inputs)
        logits_per_image = outputs.logits_per_image
        probs = logits_per_image.softmax(dim=1)

    results = {label: round(float(p) * 100, 2) for label, p in zip(candidate_labels, probs[0])}
    return results

# Интерфейс Gradio
iface = gr.Interface(
    fn=classify,
    inputs=[
        gr.Image(type="pil", label="Загрузите постер фильма"),
        gr.Textbox(label="Введите названия фильмов через запятую")
    ],
    outputs="label",
    title="Zero-Shot CLIP Movie Classifier",
    description="Сравни изображение с предложенными названиями фильмов"
)

iface.launch()