clip_filter.py

import torch
import clip
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)

# def clip_filter(model, preprocess, img_path, text, device):
def clip_filter(img_path, text):

    image = preprocess(Image.open(img_path)).unsqueeze(0).to(device)
    # text = clip.tokenize(["a diagram", "a dog", "a cat"]).to(device)
    text = clip.tokenize([text]).to(device)


    with torch.no_grad():
        image_features = model.encode_image(image)
        text_features = model.encode_text(text)
        
        # logits_per_image, logits_per_text = model(image, text)
        # probs = logits_per_image.softmax(dim=-1).cpu().numpy()

    cosine_similarity = torch.nn.functional.cosine_similarity(image_features, text_features)

    print("Cosine similarity:", cosine_similarity.item())
    return cosine_similarity.item()