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.get("/")
async def root():
    return {"message": "Gender classifier API is running. Use POST /classify/ with an image file."}

@app.post("/classify/")
async def classify_gender(image: UploadFile = File(...)):
    contents = await image.read()
    try:
        img = Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        return {"error": "Invalid image file"}

    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 = ["Female ♀", "Male ♂"]
    predictions = {labels[i]: round(probs[i], 3) for i in range(len(probs))}
    max_idx = probs.index(max(probs))

    return {
        "predictions": predictions,
        "most_likely": labels[max_idx],
        "confidence": round(probs[max_idx], 3)
    }