app.py

"""
Flask web app for FER inference.
Run: python app.py  (from the inference/ directory)

Deployment:
  Local:      weights loaded from ../models/model_weights.pth
  HF Spaces:  set HF_REPO_ID env var (e.g. "yourname/fer-weights")
              weights are downloaded from HF Model Hub at startup
"""

from __future__ import annotations

import base64
import io
import os
import sys
from pathlib import Path

import cv2
import numpy as np
from flask import Flask, jsonify, render_template, request
from flask_cors import CORS
from PIL import Image

sys.path.insert(0, str(Path(__file__).parent))
from inference import FERPredictor

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16 MB upload limit

# Allow the Vercel/GitHub Pages frontend to call /predict cross-origin
CORS(app, resources={r"/predict": {"origins": "*"}})

# Weights: download from HF Model Hub on Spaces, load locally otherwise
_HF_REPO_ID = os.getenv('HF_REPO_ID')
if _HF_REPO_ID:
    from huggingface_hub import hf_hub_download
    WEIGHTS_PATH = hf_hub_download(repo_id=_HF_REPO_ID, filename='model_weights.pth')
else:
    WEIGHTS_PATH = str(Path(__file__).parent.parent / 'models' / 'model_weights.pth')

predictor: FERPredictor | None = None


def get_predictor() -> FERPredictor:
    global predictor
    if predictor is None:
        predictor = FERPredictor(weights_path=WEIGHTS_PATH, device='auto')
    return predictor


def annotate_image(pil_img: Image.Image, face_results: list[dict]) -> str:
    """Draw bounding boxes on image, return base64-encoded JPEG."""
    img_bgr = cv2.cvtColor(np.array(pil_img.convert('RGB')), cv2.COLOR_RGB2BGR)

    EMOTION_BGR = {
        'happy':    (80, 200, 46),
        'angry':    (60, 60, 231),
        'sad':      (200, 80, 52),
        'fear':     (0, 130, 230),
        'surprise': (0, 210, 240),
        'disgust':  (150, 50, 130),
        'neutral':  (160, 160, 160),
    }

    for res in face_results:
        bbox = res.get('bbox')
        if bbox is None:
            continue
        x, y, w, h = (int(v) for v in bbox)
        emotion = res['emotion']
        conf = res['confidence']
        color = EMOTION_BGR.get(emotion, (200, 200, 200))

        cv2.rectangle(img_bgr, (x, y), (x + w, y + h), color, 2)

        label = f"{emotion}  {conf:.2f}"
        font = cv2.FONT_HERSHEY_SIMPLEX
        scale, thickness = 0.6, 2
        (tw, th), baseline = cv2.getTextSize(label, font, scale, thickness)

        ty = max(y - 6, th + 4)
        cv2.rectangle(img_bgr, (x, ty - th - 4), (x + tw + 6, ty + baseline), color, cv2.FILLED)
        cv2.putText(img_bgr, label, (x + 3, ty - 2), font, scale, (255, 255, 255), thickness, cv2.LINE_AA)

    rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
    pil_out = Image.fromarray(rgb)
    buf = io.BytesIO()
    pil_out.save(buf, format='JPEG', quality=90)
    return 'data:image/jpeg;base64,' + base64.b64encode(buf.getvalue()).decode()


def pil_to_b64(pil_img: Image.Image) -> str:
    buf = io.BytesIO()
    pil_img.convert('RGB').save(buf, format='JPEG', quality=90)
    return 'data:image/jpeg;base64,' + base64.b64encode(buf.getvalue()).decode()


@app.route('/')
def index():
    return render_template('index.html')


@app.route('/predict', methods=['POST'])
def predict():
    try:
        p = get_predictor()
        data = request.get_json(silent=True)

        if data and data.get('image'):
            # Base64 image from webcam
            b64 = data['image']
            if ',' in b64:
                b64 = b64.split(',', 1)[1]
            img_bytes = base64.b64decode(b64)
            pil_img = Image.open(io.BytesIO(img_bytes)).convert('RGB')
        elif 'file' in request.files:
            f = request.files['file']
            pil_img = Image.open(f.stream).convert('RGB')
        else:
            return jsonify({'error': 'No image provided'}), 400

        face_results = p.predict_with_face_detection(pil_img, method='mtcnn')

        # If MTCNN found no faces, fall back to whole-image prediction
        if not face_results or all(r.get('bbox') is None for r in face_results):
            result = p.predict_image(pil_img)
            result.update({'bbox': None, 'face_index': 0})
            face_results = [result]

        annotated_b64 = annotate_image(pil_img, face_results)

        faces_out = []
        for r in face_results:
            faces_out.append({
                'emotion': r['emotion'],
                'confidence': round(r['confidence'], 4),
                'probabilities': {k: round(v, 4) for k, v in r['probabilities'].items()},
                'top3': [[e, round(p_, 4)] for e, p_ in r.get('top3', [])],
                'bbox': r.get('bbox'),
            })

        return jsonify({
            'annotated_image': annotated_b64,
            'faces': faces_out,
        })

    except Exception as exc:
        import traceback
        traceback.print_exc()
        return jsonify({'error': str(exc)}), 500


if __name__ == '__main__':
    print("[INFO] Pre-loading model...")
    get_predictor()
    port = int(os.getenv('PORT', 7860))  # HF Spaces uses 7860
    app.run(host='0.0.0.0', port=port, debug=False)