app.py

from __future__ import annotations

# Block TensorFlow before any other import (crashes on machines without AVX/Rosetta)
import os
os.environ.setdefault("USE_TF",    "0")
os.environ.setdefault("USE_JAX",   "0")
os.environ.setdefault("USE_TORCH", "1")
os.environ.setdefault("TRANSFORMERS_NO_TF",  "1")
os.environ.setdefault("TRANSFORMERS_NO_JAX", "1")

import sys
from pathlib import Path

import cv2
import gradio as gr
import numpy as np

# ── model path ─────────────────────────────────────────────────────────────
MODEL_PATH = os.environ.get(
    "MODEL_PATH",
    str(Path(__file__).parent / "models" / "face_model_best.pth"),
)
MODEL_AVAILABLE = Path(MODEL_PATH).exists()

sys.path.insert(0, str(Path(__file__).parent))

# Lazy-load predictor only when model is available
_predictor = None


def get_predictor():
    global _predictor
    if _predictor is None and MODEL_AVAILABLE:
        from src.inference.predictor import Predictor
        _predictor = Predictor(model_path=MODEL_PATH)
    return _predictor


# ── emotion-only fallback (no trained model needed) ───────────────────────
_emotion_detector = None
_face_detector    = None


def get_emotion_detector():
    global _emotion_detector
    if _emotion_detector is None:
        from src.inference.emotion_detector import EmotionDetector
        _emotion_detector = EmotionDetector()
    return _emotion_detector


def get_face_detector():
    global _face_detector
    if _face_detector is None:
        from src.inference.face_detector import FaceDetector
        _face_detector = FaceDetector(confidence_threshold=0.6)
    return _face_detector


# ── inference ──────────────────────────────────────────────────────────────

def analyze(image: np.ndarray):
    """
    Entry point called by Gradio.
    Returns: (annotated_image, gender_html, age_html, emotion_html, aged_image)
    """
    try:
        return _analyze_inner(image)
    except Exception as exc:
        import traceback
        traceback.print_exc()
        err = _card(f"Error: {exc}", "Check Space logs for details")
        blank = _blank("Error")
        return blank, err, err, err, blank


def _analyze_inner(image: np.ndarray):
    if image is None:
        empty = _blank("No image received")
        return empty, "–", "–", "–", empty

    # Ensure RGB
    if image.ndim == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif image.shape[2] == 4:
        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)

    predictor = get_predictor()

    if predictor is not None:
        # Full pipeline
        results = predictor.predict_image(image)
        annotated = predictor.annotate(image)

        if not results:
            return annotated, _card("–", "No face detected"), "–", "–", _blank("No face detected")

        r            = results[0]   # use first detected face
        gender_html  = _gender_card(r["gender"], r["gender_conf"])
        age_html     = _age_card(r["age"])
        emotion_html = _emotion_card(r["emotion"], r["emotion_conf"], r["emotion_probs"])
        aged_img     = r["aged_face"]

    else:
        # Fallback: face detect + emotion only (no trained weights)
        bgr     = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
        fd      = get_face_detector()
        crops, boxes = fd.crop_faces(bgr)
        annotated = image.copy()

        if not crops:
            empty = _blank("No face detected")
            return annotated, _card("–", "No face detected"), "–", "–", empty

        for x1, y1, x2, y2 in boxes:
            cv2.rectangle(annotated, (x1, y1), (x2, y2), (52, 152, 219), 2)

        em           = get_emotion_detector()
        emotion, conf = em.top_emotion(crops[0])
        probs         = em.predict(crops[0])

        gender_html  = _card("⚠️ Model not trained yet", "Upload weights to models/")
        age_html     = _card("⚠️ Model not trained yet", "Upload weights to models/")
        emotion_html = _emotion_card(emotion, conf * 100, probs)

        from src.inference.age_progression import age_to_70
        aged_img = age_to_70(crops[0], current_age=30)

    return annotated, gender_html, age_html, emotion_html, aged_img


# ── HTML helpers ───────────────────────────────────────────────────────────

def _blank(msg: str) -> np.ndarray:
    canvas = np.ones((200, 300, 3), dtype=np.uint8) * 240
    cv2.putText(canvas, msg, (20, 110),
                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (120, 120, 120), 1)
    return canvas


def _card(title: str, subtitle: str = "") -> str:
    return f"""
<div style="background:#f8f9fa;border-radius:12px;padding:16px;text-align:center">
  <div style="font-size:1.4rem;font-weight:700;color:#2c3e50">{title}</div>
  <div style="font-size:0.85rem;color:#7f8c8d;margin-top:4px">{subtitle}</div>
</div>"""


def _gender_card(gender: str, conf: float) -> str:
    icon  = "♂" if gender == "Male" else "♀"
    color = "#3498db" if gender == "Male" else "#e74c3c"
    bar   = int(conf)
    return f"""
<div style="background:#f8f9fa;border-radius:12px;padding:16px">
  <div style="font-size:2.2rem;text-align:center;color:{color}">{icon} {gender}</div>
  <div style="background:#e0e0e0;border-radius:99px;height:8px;margin-top:10px">
    <div style="width:{bar}%;background:{color};height:8px;border-radius:99px"></div>
  </div>
  <div style="text-align:right;font-size:0.8rem;color:#7f8c8d;margin-top:4px">{conf:.1f}% confidence</div>
</div>"""


def _age_card(age: float) -> str:
    if age < 18:
        label, color = "Child / Teen", "#27ae60"
    elif age < 35:
        label, color = "Young Adult",  "#2ecc71"
    elif age < 55:
        label, color = "Middle-aged",  "#f39c12"
    else:
        label, color = "Senior",       "#e74c3c"
    return f"""
<div style="background:#f8f9fa;border-radius:12px;padding:16px;text-align:center">
  <div style="font-size:2.6rem;font-weight:800;color:{color}">{age:.0f}</div>
  <div style="font-size:0.9rem;color:#7f8c8d">years old · {label}</div>
</div>"""


def _emotion_card(emotion: str, conf: float, probs: dict) -> str:
    ICONS = {
        "Happy":    ("😊", "#f1c40f"),
        "Sad":      ("😢", "#3498db"),
        "Angry":    ("😠", "#e74c3c"),
        "Fear":     ("😨", "#9b59b6"),
        "Surprise": ("😮", "#e67e22"),
        "Disgust":  ("🤢", "#27ae60"),
        "Neutral":  ("😐", "#95a5a6"),
    }
    icon, color = ICONS.get(emotion, ("🙂", "#95a5a6"))
    bars = ""
    for lbl, prob in sorted(probs.items(), key=lambda x: -x[1]):
        w = int(prob * 100)
        ic, co = ICONS.get(lbl, ("", "#bbb"))
        bars += f"""
<div style="margin-top:5px">
  <div style="display:flex;align-items:center;gap:6px;font-size:0.78rem">
    <span>{ic}</span>
    <span style="width:70px;color:#555">{lbl}</span>
    <div style="flex:1;background:#e0e0e0;border-radius:99px;height:6px">
      <div style="width:{w}%;background:{co};height:6px;border-radius:99px"></div>
    </div>
    <span style="width:35px;text-align:right;color:#7f8c8d">{w}%</span>
  </div>
</div>"""

    return f"""
<div style="background:#f8f9fa;border-radius:12px;padding:16px">
  <div style="font-size:1.8rem;text-align:center">{icon} <span style="color:{color};font-weight:700">{emotion}</span></div>
  <div style="font-size:0.82rem;color:#7f8c8d;text-align:center;margin-bottom:8px">{conf:.1f}% confidence</div>
  {bars}
</div>"""


# ── CSS ────────────────────────────────────────────────────────────────────

CSS = """
/* ---------- globals ---------- */
body, .gradio-container { font-family: 'Inter', system-ui, sans-serif !important; }
.gradio-container { max-width: 1100px !important; margin: 0 auto !important; }

/* ---------- hero ---------- */
.hero {
  background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%);
  border-radius: 20px;
  padding: 40px 30px 30px;
  margin-bottom: 24px;
  text-align: center;
  color: white;
}
.hero h1 {
  font-size: clamp(2rem, 5vw, 3.2rem);
  font-weight: 800;
  letter-spacing: -1px;
  margin: 0 0 10px;
  background: linear-gradient(90deg, #00d2ff, #a8edea);
  -webkit-background-clip: text;
  -webkit-text-fill-color: transparent;
}
.hero p {
  font-size: clamp(0.9rem, 2.5vw, 1.05rem);
  color: #a8c0d6;
  max-width: 620px;
  margin: 0 auto 18px;
  line-height: 1.6;
}
.badge {
  display: inline-block;
  background: rgba(255,255,255,0.12);
  border: 1px solid rgba(255,255,255,0.2);
  border-radius: 99px;
  padding: 4px 14px;
  font-size: 0.78rem;
  color: #a8edea;
  margin: 0 4px;
}

/* ---------- panels ---------- */
.panel { background: white; border-radius: 16px; padding: 20px; box-shadow: 0 2px 12px rgba(0,0,0,0.07); }
.section-label {
  font-size: 0.72rem;
  font-weight: 700;
  letter-spacing: 1.2px;
  text-transform: uppercase;
  color: #95a5a6;
  margin-bottom: 8px;
}

/* ---------- webcam button ---------- */
.analyze-btn {
  background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
  color: white !important;
  font-weight: 700 !important;
  font-size: 1rem !important;
  border-radius: 12px !important;
  padding: 14px !important;
  border: none !important;
  width: 100% !important;
  cursor: pointer !important;
  transition: opacity 0.2s !important;
}
.analyze-btn:hover { opacity: 0.88 !important; }

/* ---------- aged-face label ---------- */
.aged-label {
  font-size: 0.82rem;
  color: #7f8c8d;
  text-align: center;
  margin-top: 6px;
}

/* ---------- footer ---------- */
.footer {
  text-align: center;
  font-size: 0.78rem;
  color: #bbb;
  margin-top: 16px;
  padding-bottom: 10px;
}

/* ---------- responsive ---------- */
@media (max-width: 700px) {
  .hero { padding: 28px 16px 20px; }
  .gradio-container { padding: 8px !important; }
}
"""

HERO_HTML = """
<div class="hero">
  <h1>🧠 FaceInsight AI</h1>
  <p>
    Point your camera at a face — or upload a photo — and instantly see
    <strong>gender</strong>, <strong>age</strong>, <strong>emotion</strong>,
    and a preview of <strong>how the person will look at age&nbsp;70</strong>.
  </p>
  <span class="badge">✦ Real-time webcam</span>
  <span class="badge">✦ Works on mobile & desktop</span>
  <span class="badge">✦ No data stored</span>
</div>
"""

HOW_HTML = """
<div style="background:#eaf4ff;border-radius:12px;padding:14px 18px;font-size:0.86rem;color:#2c3e50;line-height:1.7">
  <strong>How to use:</strong><br>
  1️⃣  Click <em>Allow</em> when your browser asks for camera access.<br>
  2️⃣  Align your face in the frame — good lighting helps!<br>
  3️⃣  Press <strong>Analyze Face</strong> or upload a photo from your gallery.<br>
  4️⃣  Results appear instantly on the right. The side panel shows a simulated portrait of you at 70.
</div>
"""

# ── build Gradio app ───────────────────────────────────────────────────────

def build_app() -> gr.Blocks:
    with gr.Blocks(title="FaceInsight_AI") as demo:

        gr.HTML(HERO_HTML)

        with gr.Row(equal_height=False):

            # ── LEFT COLUMN ───────────────────────────────────────────────
            with gr.Column(scale=5, min_width=280):
                gr.HTML(HOW_HTML)

                image_input = gr.Image(
                    sources    = ["webcam", "upload"],
                    type       = "numpy",
                    label      = "📷  Camera / Upload",
                    height     = 360,
                    show_label = True,
                )

                analyze_btn = gr.Button(
                    "🔍  Analyze Face",
                    elem_classes = ["analyze-btn"],
                    variant      = "primary",
                )

            # ── RIGHT COLUMN ──────────────────────────────────────────────
            with gr.Column(scale=7, min_width=320):

                annotated_out = gr.Image(
                    label  = "🎯  Detected faces",
                    height = 320,
                )

                with gr.Row():
                    gender_out  = gr.HTML(label="Gender")
                    age_out     = gr.HTML(label="Age")

                emotion_out = gr.HTML(label="Emotion")

        # ── AGED FACE (full width below) ──────────────────────────────────
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML("""
<div style="text-align:center;margin:20px 0 8px">
  <span style="font-size:1.1rem;font-weight:700;color:#2c3e50">🕰️ Simulated portrait — age 70</span><br>
  <span style="font-size:0.82rem;color:#95a5a6">This is an artistic simulation, not medical or forensic analysis.</span>
</div>""")
                aged_out = gr.Image(
                    label  = "You at 70",
                    height = 300,
                )

        # ── EXAMPLES ──────────────────────────────────────────────────────
        gr.Examples(
            examples        = [["examples/sample1.jpg"]],
            inputs          = [image_input],
            outputs         = [annotated_out, gender_out, age_out, emotion_out, aged_out],
            fn              = analyze,
            cache_examples  = False,
            label           = "Try a sample image",
        ) if Path("examples/sample1.jpg").exists() else None

        # ── wire up ───────────────────────────────────────────────────────
        analyze_btn.click(
            fn      = analyze,
            inputs  = [image_input],
            outputs = [annotated_out, gender_out, age_out, emotion_out, aged_out],
        )

        # Run on webcam changes only when image is not None
        image_input.change(
            fn      = lambda img: analyze(img) if img is not None else (None,)*5,
            inputs  = [image_input],
            outputs = [annotated_out, gender_out, age_out, emotion_out, aged_out],
        )

        gr.HTML("""
<div class="footer">
  FaceInsight AI · Trained on UTKFace (White/Black = US · Indian) ·
  Emotion: FER-2013 · Age-at-70 is an artistic effect only ·
  No images are stored or transmitted.
</div>""")

    return demo


if __name__ == "__main__":
    app = build_app()
    app.launch(
        server_name = "0.0.0.0",
        server_port = int(os.environ.get("PORT", 7860)),
        css         = CSS,
        share       = True,    # creates a public gradio.live tunnel URL
    )