import torch
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from pathlib import Path

# Пути к папкам
image_dir = Path("images")
text_dir = Path("texts")

# Модель
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

results = []

for image_path in image_dir.glob("*"):
    if image_path.suffix.lower() not in [".png", ".jpg", ".jpeg"]:
        continue

    # имя без расширения
    name = image_path.stem
    text_path = text_dir / name

    if not text_path.exists():
        print(f"Нет текста для {name}")
        continue

    # загрузка
    image = Image.open(image_path)
    text = text_path.read_text(encoding="utf-8").strip()

    inputs = processor(
        text=[text],
        images=image,
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=77
    )

    with torch.no_grad():
        outputs = model(**inputs)

        image_embeds = outputs.image_embeds
        text_embeds = outputs.text_embeds

        image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)

        score = (image_embeds @ text_embeds.T).item()

    if score > 0.27:
        score *= 300
    elif 0.2 < score <= 0.27:
        score *= 200

    print(f"{name}: {score:.4f}%")
    results.append((name, score))

# например, средний скор
if results:
    avg_score = sum(s for _, s in results) / len(results)
    print(f"\nAverage score: {avg_score:.4f}%")