modules/embeddings.py

'''
import torch
from torchvision import models, transforms
from PIL import Image
import requests
from io import BytesIO

# Load ResNet50 model
model = models.resnet50(pretrained=True)
model.eval()

preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
])

def get_embedding(image: Image.Image):
    img_tensor = preprocess(image).unsqueeze(0)
    with torch.no_grad():
        embedding = model(img_tensor).squeeze().numpy()
    return embedding

def load_image_from_url(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content)).convert("RGB")
    return img
'''   
    
import torch
import open_clip
from PIL import Image
import requests
from io import BytesIO
from torchvision import transforms

# Load CLIP model (ViT-B-32 is a good default)
model, _, preprocess = open_clip.create_model_and_transforms(
    'ViT-B-32', pretrained='openai'
)
model.eval()

# Define a function to compute CLIP embeddings
def get_embedding(image: Image.Image):
    # CLIP expects images in a certain format:
    image_input = preprocess(image).unsqueeze(0)  # (1, 3, H, W)
    with torch.no_grad():
        image_features = model.encode_image(image_input)
    # Normalize to unit vector for cosine similarity
    image_features = image_features / image_features.norm(dim=-1, keepdim=True)
    return image_features.squeeze().cpu().numpy()

# Helper to load image from URL
def load_image_from_url(url):
    response = requests.get(url)
    img = Image.open(BytesIO(response.content)).convert("RGB")
    return img