Update app.py

app.py

@@ -1,35 +1,28 @@
 import os
-# Disable GPU for TensorFlow / Keras BEFORE importing anything
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
-
 import time
 import cv2
 import numpy as np
 import torch
-from flask import Flask, request, jsonify
-from ultralytics import YOLO
 from PIL import Image as PILImage
-from datetime import datetime, timedelta
 import gc
+from datetime import datetime, timedelta
 
+from ultralytics import YOLO
 from keras_facenet import FaceNet
 from transformers import pipeline
-
-# -----------------------------
-# Flask Setup
-# -----------------------------
-app = Flask(__name__)
+import gradio as gr
 
 # -----------------------------
 # Device Setup
 # -----------------------------
+os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # Disable GPU
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Using device: {DEVICE}")
 
 # -----------------------------
 # Load YOLOv8 Face Model
 # -----------------------------
-MODEL_PATH = "yolov8n-face.pt"   # put this file in your repo root
+MODEL_PATH = "yolov8n-face.pt"  # make sure this is in your repo
 if not os.path.exists(MODEL_PATH):
     raise FileNotFoundError(f"Model file not found at {MODEL_PATH}")
 
@@ -40,18 +33,18 @@ print("YOLOv8 loaded")
 # -----------------------------
 # Load FaceNet Embedder
 # -----------------------------
-print("Loading FaceNet (Google) model...")
-embedder = FaceNet()  # CPU mode
+print("Loading FaceNet model...")
+embedder = FaceNet()
 print("FaceNet loaded")
 
 # -----------------------------
 # Load HuggingFace Age & Gender Models
 # -----------------------------
-print("Loading HuggingFace models...")
+print("Loading HuggingFace Age & Gender models...")
 age_model = pipeline(
     "image-classification",
     model="prithivMLmods/Age-Classification-SigLIP2",
-    device=-1  # CPU mode
+    device=-1
 )
 gender_model = pipeline(
     "image-classification",
@@ -67,127 +60,112 @@ FACE_DB = []
 NEXT_ID = 1
 
 # -----------------------------
-# GPU Cleaner
+# Utilities
 # -----------------------------
 def clean_gpu():
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
     gc.collect()
 
-# -----------------------------
-# Cosine Similarity
-# -----------------------------
 def cosine_similarity(a, b):
     return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
 
 # -----------------------------
-# Upload API Endpoint
+# Core Function
 # -----------------------------
-@app.route("/upload", methods=["POST"])
-def upload():
+def process_image(image: PILImage):
     global NEXT_ID, FACE_DB
     start_time = time.time()
-
-    try:
-        # Decode image
-        jpg_bytes = request.data
-        np_arr = np.frombuffer(jpg_bytes, np.uint8)
-        img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
-        if img is None:
-            return jsonify({"status": "error", "message": "Invalid image"}), 400
-
-        rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-
-        # Detect faces using YOLOv8
-        results = face_model(rgb_img, verbose=False)
-        boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
-
-        now = datetime.now()
-        # Remove old entries (> 1 hour)
-        FACE_DB = [f for f in FACE_DB if now - f["time"] <= timedelta(hours=1)]
-
-        faces = []
-        for (x1, y1, x2, y2) in boxes:
-            face_crop = rgb_img[y1:y2, x1:x2]
-            if face_crop.size == 0:
-                continue
-
-            # Get embedding from FaceNet
-            face_embedding = embedder.embeddings([face_crop])[0]
-
-            assigned_id = None
-            age_pred, gender_pred = "unknown", "unknown"
-
-            # Compare with known embeddings
-            if FACE_DB:
-                similarities = [cosine_similarity(face_embedding, entry["embedding"]) for entry in FACE_DB]
-                best_match_index = int(np.argmax(similarities))
-                best_score = similarities[best_match_index]
-
-                if best_score > 0.6:  # same person threshold
-                    assigned_id = FACE_DB[best_match_index]["id"]
-                    FACE_DB[best_match_index]["time"] = now
-                    FACE_DB[best_match_index]["seen_count"] += 1
-                    age_pred = FACE_DB[best_match_index]["age"]
-                    gender_pred = FACE_DB[best_match_index]["gender"]
-
-            # If new person
-            if assigned_id is None:
-                assigned_id = NEXT_ID
-                face_pil = PILImage.fromarray(face_crop)
-
-                try:
-                    age_pred = age_model(face_pil)[0]["label"]
-                    gender_pred = gender_model(face_pil)[0]["label"]
-                except Exception:
-                    age_pred, gender_pred = "unknown", "unknown"
-
-                FACE_DB.append({
-                    "id": assigned_id,
-                    "embedding": face_embedding,
-                    "time": now,
-                    "seen_count": 1,
-                    "age": age_pred,
-                    "gender": gender_pred
-                })
-                NEXT_ID += 1
-
-            faces.append({
+    rgb_img = np.array(image)
+
+    # Detect faces
+    results = face_model(rgb_img, verbose=False)
+    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
+
+    now = datetime.now()
+    # Remove old entries (>1 hour)
+    FACE_DB = [f for f in FACE_DB if now - f["time"] <= timedelta(hours=1)]
+
+    faces = []
+    for (x1, y1, x2, y2) in boxes:
+        face_crop = rgb_img[y1:y2, x1:x2]
+        if face_crop.size == 0:
+            continue
+
+        face_embedding = embedder.embeddings([face_crop])[0]
+        assigned_id = None
+        age_pred, gender_pred = "unknown", "unknown"
+
+        # Compare with DB
+        if FACE_DB:
+            similarities = [cosine_similarity(face_embedding, entry["embedding"]) for entry in FACE_DB]
+            best_idx = int(np.argmax(similarities))
+            best_score = similarities[best_idx]
+            if best_score > 0.6:
+                assigned_id = FACE_DB[best_idx]["id"]
+                FACE_DB[best_idx]["time"] = now
+                FACE_DB[best_idx]["seen_count"] += 1
+                age_pred = FACE_DB[best_idx]["age"]
+                gender_pred = FACE_DB[best_idx]["gender"]
+
+        # New face
+        if assigned_id is None:
+            assigned_id = NEXT_ID
+            face_pil = PILImage.fromarray(face_crop)
+            try:
+                age_pred = age_model(face_pil)[0]["label"]
+                gender_pred = gender_model(face_pil)[0]["label"]
+            except Exception:
+                age_pred, gender_pred = "unknown", "unknown"
+
+            FACE_DB.append({
                 "id": assigned_id,
+                "embedding": face_embedding,
+                "time": now,
+                "seen_count": 1,
                 "age": age_pred,
-                "gender": gender_pred,
-                "box": [int(x1), int(y1), int(x2), int(y2)]
+                "gender": gender_pred
             })
+            NEXT_ID += 1
 
-        total_time = round(time.time() - start_time, 3)
-        clean_gpu()
-
-        summary = [
-            {
-                "id": entry["id"],
-                "seen_count": entry["seen_count"],
-                "age": entry["age"],
-                "gender": entry["gender"]
-            }
-            for entry in FACE_DB
-        ]
-
-        return jsonify({
-            "status": "ok",
-            "faces": faces,
-            "face_count": len(faces),
-            "processing_time_sec": total_time,
-            "active_faces_last_hour": len(FACE_DB),
-            "seen_summary_last_hour": summary
+        faces.append({
+            "id": assigned_id,
+            "age": age_pred,
+            "gender": gender_pred,
+            "box": [int(x1), int(y1), int(x2), int(y2)]
         })
 
-    except Exception as e:
-        return jsonify({"status": "error", "message": str(e)}), 500
-
+    total_time = round(time.time() - start_time, 3)
+    clean_gpu()
+
+    summary = [
+        {
+            "id": entry["id"],
+            "seen_count": entry["seen_count"],
+            "age": entry["age"],
+            "gender": entry["gender"]
+        } for entry in FACE_DB
+    ]
+
+    return {
+        "status": "ok",
+        "faces": faces,
+        "face_count": len(faces),
+        "processing_time_sec": total_time,
+        "active_faces_last_hour": len(FACE_DB),
+        "seen_summary_last_hour": summary
+    }
+
+# -----------------------------
+# Gradio Interface
+# -----------------------------
+demo = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil"),
+    outputs="json",
+    title="Face Recognition + Age & Gender",
+    description="YOLOv8 + FaceNet + HuggingFace Age/Gender"
+)
 
-# -----------------------------
-# Run Server
-# -----------------------------
 if __name__ == "__main__":
-    # Hugging Face Spaces expects host 0.0.0.0 and port 7860
-    app.run(host="0.0.0.0", port=7860, debug=False, use_reloader=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860)