app/model.py

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

MODEL_NAME = "openai/clip-vit-base-patch16"

# Load model and processor only once at startup
processor = CLIPProcessor.from_pretrained(MODEL_NAME)
model = CLIPModel.from_pretrained(MODEL_NAME)

# Define the categories to classify into
CATEGORIES = ["food", "fitness", "healthcare"]

def analyze_image(image: Image.Image):
    # Preprocess input
    inputs = processor(
        text=CATEGORIES,
        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).cpu().numpy()[0]
        best_idx = probs.argmax()
        return {
            "category": CATEGORIES[best_idx],
            "confidence": round(float(probs[best_idx]), 4)
        }