app.py

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)