Upload 2 files

app.py

@@ -0,0 +1,151 @@
+import gradio as gr
+from transformers import ViTImageProcessor, ViTForImageClassification
+from transformers import AutoImageProcessor, SiglipForImageClassification
+from PIL import Image
+import torch
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
+import io
+import logging
+
+# ----------------- LOGGER SETUP -----------------
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logger = logging.getLogger("face-analysis")
+
+# ----------------- LOAD MODELS -----------------
+# Emotion model
+emotion_processor = ViTImageProcessor.from_pretrained("abhilash88/face-emotion-detection")
+emotion_model = ViTForImageClassification.from_pretrained("abhilash88/face-emotion-detection")
+
+# Age model
+age_model_name = "prithivMLmods/facial-age-detection"
+age_model = SiglipForImageClassification.from_pretrained(age_model_name)
+age_processor = AutoImageProcessor.from_pretrained(age_model_name)
+
+# Emotion classes
+emotions = ["Angry", "Disgust", "Fear", "Happy", "Sad", "Surprise", "Neutral"]
+
+# Age labels
+id2label = {
+    "0": "age 01-10",
+    "1": "age 11-20",
+    "2": "age 21-30",
+    "3": "age 31-40",
+    "4": "age 41-55",
+    "5": "age 56-65",
+    "6": "age 66-80",
+    "7": "age 80+"
+}
+
+# ----------------- PREDICT FUNCTIONS -----------------
+def predict_emotion(image: Image.Image):
+    try:
+        inputs = emotion_processor(image.convert("RGB"), return_tensors="pt")
+        with torch.no_grad():
+            outputs = emotion_model(**inputs)
+            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]
+        idx = torch.argmax(probs).item()
+
+        result = {
+            "predicted_emotion": emotions[idx],
+            "confidence": round(probs[idx].item(), 4),
+            "all_confidences": {emotions[i]: float(probs[i]) for i in range(len(emotions))}
+        }
+
+        logger.info(f"Predicted Emotion: {result['predicted_emotion']} (Confidence: {result['confidence']})")
+        return result
+    except Exception as e:
+        logger.error(f"Emotion prediction error: {e}")
+        return {"error": str(e)}
+
+def predict_age(image: Image.Image):
+    try:
+        inputs = age_processor(images=image.convert("RGB"), return_tensors="pt")
+        with torch.no_grad():
+            outputs = age_model(**inputs)
+            probs = torch.nn.functional.softmax(outputs.logits, dim=1).squeeze().tolist()
+        prediction = {id2label[str(i)]: round(probs[i], 3) for i in range(len(probs))}
+        idx = int(torch.argmax(torch.tensor(probs)))
+
+        result = {
+            "predicted_age": id2label[str(idx)],
+            "confidence": round(probs[idx], 4),
+            "all_confidences": prediction
+        }
+
+        logger.info(f"Predicted Age Group: {result['predicted_age']} (Confidence: {result['confidence']})")
+        return result
+    except Exception as e:
+        logger.error(f"Age prediction error: {e}")
+        return {"error": str(e)}
+
+# ----------------- FASTAPI APP -----------------
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
+
+@app.post("/predict")
+async def predict(file: UploadFile = File(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+
+        emotion_result = predict_emotion(image)
+        age_result = predict_age(image)
+
+        logger.info(f"API Response -> Emotion: {emotion_result.get('predicted_emotion')} | Age: {age_result.get('predicted_age')}")
+
+        return JSONResponse(content={
+            "emotion": emotion_result,
+            "age": age_result
+        })
+    except Exception as e:
+        logger.error(f"API Error: {e}")
+        return JSONResponse(content={"error": str(e)}, status_code=500)
+
+# ----------------- GRADIO DEMO -----------------
+def gradio_wrapper(image):
+    emotion_result = predict_emotion(image)
+    age_result = predict_age(image)
+
+    if "error" in emotion_result or "error" in age_result:
+        return "Error", {}, "Error", {}
+
+    return (
+        f"{emotion_result['predicted_emotion']} ({emotion_result['confidence']:.2f})",
+        emotion_result["all_confidences"],
+        f"{age_result['predicted_age']} ({age_result['confidence']:.2f})",
+        age_result["all_confidences"]
+    )
+
+demo = gr.Interface(
+    fn=gradio_wrapper,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Label(num_top_classes=1, label="Top Emotion"),
+        gr.Label(label="Emotion Probabilities"),
+        gr.Label(num_top_classes=1, label="Top Age Group"),
+        gr.Label(label="Age Probabilities"),
+    ],
+    title="Face Emotion + Age Detection",
+    description="Upload a face image and detect both emotion (Angry, Happy, etc.) and estimated age group (01–10, 11–20, ... 80+)."
+)
+
+# Mount Gradio at /gradio
+app = gr.mount_gradio_app(app, demo, path="/gradio")
+
+# ----------------- RUN -----------------
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -0,0 +1,4 @@
+torch
+transformers
+gradio
+pillow