Update app.py

app.py

@@ -1,141 +1,193 @@
-import os
-import time
-import cv2
-import numpy as np
-import torch
-from ultralytics import YOLO
-from PIL import Image as PILImage
-from keras_facenet import FaceNet
-from transformers import pipeline
-import gradio as gr
-from datetime import datetime, timedelta
-import gc
-
-# -----------------------------
-# Device Setup
-# -----------------------------
-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 Space repository
-face_model = YOLO(MODEL_PATH).to(DEVICE)
-
-# -----------------------------
-# Load FaceNet Embedder
-# -----------------------------
-embedder = FaceNet()  # CPU mode
-
-# -----------------------------
-# Load HuggingFace Age & Gender Models
-# -----------------------------
-age_model = pipeline(
-    "image-classification",
-    model="prithivMLmods/Age-Classification-SigLIP2",
-    device=-1
-)
-gender_model = pipeline(
-    "image-classification",
-    model="dima806/fairface_gender_image_detection",
-    device=-1
-)
-
-# -----------------------------
-# Face DB
-# -----------------------------
-FACE_DB = []
-NEXT_ID = 1
-
-def clean_gpu():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    gc.collect()
-
-def cosine_similarity(a, b):
-    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
-
-# -----------------------------
-# Main Inference Function
-# -----------------------------
-def process_image(image):
-    global NEXT_ID, FACE_DB
-    start_time = time.time()
-
-    rgb_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-    # Detect faces
-    results = face_model(rgb_img, verbose=False)
-    boxes = results[0].boxes.xyxy.cpu().numpy().astype(int)
-
-    now = datetime.now()
-    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_match_index = int(np.argmax(similarities))
-            best_score = similarities[best_match_index]
-
-            if best_score > 0.6:
-                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 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({
-            "id": assigned_id,
-            "age": age_pred,
-            "gender": gender_pred,
-            "box": [int(x1), int(y1), int(x2), int(y2)]
-        })
-
-    total_time = round(time.time() - start_time, 3)
-    clean_gpu()
-
-    return {"faces": faces, "processing_time_sec": total_time}
-
-# -----------------------------
-# Gradio UI
-# -----------------------------
-iface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="numpy"),
-    outputs="json",
-    title="Face Detection + Age/Gender"
-)
-
-if __name__ == "__main__":
-    iface.launch()
+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 keras_facenet import FaceNet
+from transformers import pipeline
+
+# -----------------------------
+# Flask Setup
+# -----------------------------
+app = Flask(__name__)
+
+# -----------------------------
+# Device Setup
+# -----------------------------
+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
+if not os.path.exists(MODEL_PATH):
+    raise FileNotFoundError(f"Model file not found at {MODEL_PATH}")
+
+print("Loading YOLOv8 face detector...")
+face_model = YOLO(MODEL_PATH).to(DEVICE)
+print("YOLOv8 loaded")
+
+# -----------------------------
+# Load FaceNet Embedder
+# -----------------------------
+print("Loading FaceNet (Google) model...")
+embedder = FaceNet()  # CPU mode
+print("FaceNet loaded")
+
+# -----------------------------
+# Load HuggingFace Age & Gender Models
+# -----------------------------
+print("Loading HuggingFace models...")
+age_model = pipeline(
+    "image-classification",
+    model="prithivMLmods/Age-Classification-SigLIP2",
+    device=-1  # CPU mode
+)
+gender_model = pipeline(
+    "image-classification",
+    model="dima806/fairface_gender_image_detection",
+    device=-1
+)
+print("Age & Gender models loaded")
+
+# -----------------------------
+# Face DB
+# -----------------------------
+FACE_DB = []
+NEXT_ID = 1
+
+# -----------------------------
+# GPU Cleaner
+# -----------------------------
+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
+# -----------------------------
+@app.route("/upload", methods=["POST"])
+def upload():
+    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({
+                "id": assigned_id,
+                "age": age_pred,
+                "gender": gender_pred,
+                "box": [int(x1), int(y1), int(x2), int(y2)]
+            })
+
+        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
+        })
+
+    except Exception as e:
+        return jsonify({"status": "error", "message": str(e)}), 500
+
+
+# -----------------------------
+# 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)