router/ai/scripts/clip_test.py

from PIL import Image
import requests
import torch

from transformers import CLIPProcessor, CLIPModel


def main() -> None:
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

    # Default test image; change this URL or switch to local file path if needed.
    url = "http://images.cocodataset.org/val2017/000000039769.jpg"
    image = Image.open(requests.get(url, stream=True, timeout=30).raw).convert("RGB")

    labels = ["a photo of a cat", "a photo of a dog"]
    inputs = processor(text=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)[0]

    for label, score in zip(labels, probs.tolist()):
        print(f"{label}: {score:.6f}")

    best_idx = int(probs.argmax().item())
    print(f"\nPredicted label: {labels[best_idx]}")


if __name__ == "__main__":
    main()