app.py

import os
os.environ["HF_HOME"] = "/tmp/huggingface"

from fastapi import FastAPI, UploadFile, File
from transformers import SiglipForImageClassification, AutoImageProcessor
from PIL import Image
import torch
import torch.nn.functional as F
import io

app = FastAPI()

model_name = "prithivMLmods/Gender-Classifier-Mini"
model = SiglipForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)

@app.post("/classify/")
async def classify_gender(image: UploadFile = File(...)):
    contents = await image.read()
    img = Image.open(io.BytesIO(contents)).convert("RGB")
    inputs = processor(images=img, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = F.softmax(logits, dim=1).squeeze().tolist()

    labels = {"0": "Female ♀", "1": "Male ♂"}
    predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}

    return predictions