Update app.py

app.py

@@ -14,10 +14,18 @@ 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()
-    img = Image.open(io.BytesIO(contents)).convert("RGB")
+    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():
@@ -25,7 +33,12 @@ async def classify_gender(image: UploadFile = File(...)):
         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))}
+    labels = ["Female ♀", "Male ♂"]
+    predictions = {labels[i]: round(probs[i], 3) for i in range(len(probs))}
+    max_idx = probs.index(max(probs))
 
-    return predictions
+    return {
+        "predictions": predictions,
+        "most_likely": labels[max_idx],
+        "confidence": round(probs[max_idx], 3)
+    }