utils/inference.py

# import requests

# API_URL = "https://api-inference.huggingface.co/models/trpakov/vit-face-expression"
# headers = {"Authorization": "Bearer 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"}


# def query(filename):
#     with open(filename, "rb") as f:
#         data = f.read()
#     response = requests.post(API_URL, headers=headers, data=data)
#     return response.json()


from PIL import Image
from transformers import CLIPProcessor, CLIPModel

model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")


def query(filename):
    image = Image.open(filename)
    inputs = processor(
        text=["Happy", "Sad", "Surprised", "Angry", "Disgusted", "Neutral", "Fearful"],
        images=image,
        return_tensors="pt",
        padding=True,
    )
    outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image
    probs = logits_per_image.softmax(dim=1)
    output = [{"label": label, "score": float(score)} for label, score in zip(["Happy", "Sad", "Surprised", "Angry", "Disgusted", "Neutral", "Fearful"], probs[0])]
    return output