app.py

"""
Gradio frontend — Flood Detection + ZoeDepth depth estimation.
Improved UI: dark theme, styled cards, progress steps, metric badges.
"""
import os
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
os.environ["TF_XLA_FLAGS"] = "--tf_xla_auto_jit=0"
import gradio as gr
import spaces
import numpy as np
from PIL import Image
import sys
import asyncio

# Suppress harmless asyncio garbage collection bug in Python 3.13
if hasattr(asyncio, "base_events") and hasattr(asyncio.base_events, "BaseEventLoop"):
    _original_del = asyncio.base_events.BaseEventLoop.__del__
    def _safe_del(self):
        try:
            _original_del(self)
        except Exception:
            pass
    asyncio.base_events.BaseEventLoop.__del__ = _safe_del

import base64
from gradio.themes.utils.colors import Color

sys.path.insert(0, os.path.dirname(__file__))
from app.model_utils import run_pipeline

def get_base64_image(path):
    try:
        with open(path, "rb") as img_file:
            return base64.b64encode(img_file.read()).decode('utf-8')
    except Exception:
        return ""

logo_base64 = get_base64_image(os.path.join(os.path.dirname(__file__), "clipart1553592.png"))
logo_src = f"data:image/png;base64,{logo_base64}"

# ── Risk config ────────────────────────────────────────────────────────────────
RISK_CFG = {
    "Low":      {"colour": "#27ae60", "bg": "#eafaf1", "bar": 20},
    "Moderate": {"colour": "#f39c12", "bg": "#fef9e7", "bar": 50},
    "High":     {"colour": "#e67e22", "bg": "#fdf2e9", "bar": 75},
    "Critical": {"colour": "#e74c3c", "bg": "#fdedec", "bar": 100},
}

CSS = """
@import url('https://fonts.googleapis.com/css2?family=Geist:wght@400;500;600;700;800&display=swap');

:root, .dark {
    --primary-50: #f0f7fb !important;
    --primary-100: #e1f0f7 !important;
    --primary-200: #c4e1ef !important;
    --primary-300: #a6d2e8 !important;
    --primary-400: #89c3e0 !important;
    --primary-500: #43a0d6 !important;
    --primary-600: #3680ab !important;
    --primary-700: #286080 !important;
    --primary-800: #1b4056 !important;
    --primary-900: #0d202b !important;
    --primary-950: #071015 !important;
}

/* ── light theme overrides (warm off-white) ── */
.light {
    --block-background-fill: #f7f3eb !important;
    --background-fill-primary: #f7f3eb !important;
    --background-fill-secondary: #eee9dd !important;
    --border-color-primary: #dfd8c9 !important;
    --body-text-color: #1f2937 !important;
    --body-text-color-subdued: #4b5563 !important;
    --block-label-text-color: #1f2937 !important;
    --block-info-text-color: #4b5563 !important;
}
.light .upload-container span, 
.light .upload-container p,
.light .upload-container div {
    color: #4b5563 !important;
}

/* ── page ── */
body, .gradio-container { 
    font-family: 'Geist', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important; 
}
.light body, .light .gradio-container { background-color: #f7f3eb !important; }

/* ── theme toggle switch ── */
.theme-switch {
  position: absolute; top: 32px; right: 36px;
  display: inline-block; width: 44px; height: 24px; z-index: 10;
}
.theme-switch input { opacity: 0; width: 0; height: 0; }
.theme-slider {
  position: absolute; cursor: pointer; top: 0; left: 0; right: 0; bottom: 0;
  background-color: #cbd5e1; transition: .4s; border-radius: 24px;
}
.theme-slider:before {
  position: absolute; content: ""; height: 18px; width: 18px; left: 3px; bottom: 3px;
  background-color: #f1f5f9; transition: .4s; border-radius: 50%; box-shadow: 0 1px 3px rgba(0,0,0,0.2);
}
input:checked + .theme-slider { background-color: #3ba3e8; }
input:checked + .theme-slider:before { transform: translateX(20px); background-color: #f8fafc; }

#header-band {
    background: var(--block-background-fill);
    border-radius: 12px; 
    position: relative;
    padding: 32px 36px 24px; 
    margin-bottom: 24px;
    box-shadow: var(--block-shadow);
    border: 1px solid var(--border-color-primary);
}
#header-band h1 { color: var(--body-text-color); font-size: 2.2em; margin: 0 0 8px; font-weight: 800; letter-spacing: -0.03em; }
#header-band p  { color: var(--body-text-color-subdued); margin: 0; font-size: 1.05em; font-weight: 400; }

/* ── metric pill row ── */
.metric-pill {
    display: inline-block; background: var(--background-fill-secondary);
    border: 1px solid var(--border-color-primary); border-radius: 8px;
    padding: 6px 16px; margin: 6px 6px 6px 0; font-size: 0.9em; color: var(--body-text-color-subdued); font-weight: 500;
}
.metric-pill strong { color: var(--body-text-color); font-weight: 700; }

/* ── upload + button panel ── */
#left-panel { background: var(--block-background-fill); border-radius: 12px; padding: 20px; border: 1px solid var(--border-color-primary); box-shadow: var(--block-shadow); }

/* ── analyse button ── */
#analyse-btn {
    background: var(--primary-500) !important;
    color: white !important; border: none !important;
    border-radius: 8px !important; font-size: 1.1em !important; font-weight: 600 !important;
    padding: 14px !important; margin-top: 12px !important;
    box-shadow: 0 4px 14px rgba(0,0,0,0.1) !important;
    transition: all .2s ease !important;
}
#analyse-btn:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 8px 25px rgba(0,0,0,0.15) !important;
    background: var(--primary-600) !important;
}

/* ── section labels ── */
.section-label {
    font-size: 0.85em; font-weight: 700; letter-spacing: .1em;
    text-transform: uppercase; color: var(--body-text-color-subdued); margin: 24px 0 10px;
}

/* ── image cards ── */
.image-card { background: var(--block-background-fill) !important; border-radius: 12px !important;
              border: 1px solid var(--border-color-primary) !important; overflow: hidden; box-shadow: var(--block-shadow) !important; }

/* ── risk card ── */
#risk-card { border-radius: 12px; overflow: hidden; background: var(--block-background-fill); box-shadow: var(--block-shadow); border: 1px solid var(--border-color-primary); }

/* ── how-it-works table ── */
.how-table { width:100%; border-collapse:collapse; font-size:0.95em; }
.how-table th { background: var(--background-fill-secondary); color: var(--body-text-color); padding:14px 16px; text-align:left; border-bottom: 2px solid var(--border-color-primary); font-weight: 600; }
.how-table td { padding:14px 16px; border-bottom:1px solid var(--border-color-primary); color: var(--body-text-color); }
.how-table tr:last-child td { border-bottom: none; }
"""

# ── helpers ───────────────────────────────────────────────────────────────────
def _score_bar(score: float, colour: str) -> str:
    """Animated CSS progress bar."""
    return f"""
    <div style="background:var(--background-fill-secondary); border-radius:99px; height:12px; margin:12px 0 20px; overflow:hidden;">
      <div style="width:{score}%; background:linear-gradient(90deg,{colour}bb,{colour});
                  height:100%; border-radius:99px;
                  transition: width 1s ease;"></div>
    </div>"""

def _stat_row(label, value, colour) -> str:
    return f"""
    <div style="display:flex; justify-content:space-between; align-items:center;
                padding:12px 0; border-bottom:1px solid var(--border-color-primary);">
      <span style="color:var(--body-text-color-subdued); font-weight: 500;">{label}</span>
      <span style="font-weight:700; color:{colour};">{value}</span>
    </div>"""

def build_risk_html(risk: dict, depth_info: dict) -> str:
    level  = risk["risk_level"]
    cfg    = RISK_CFG.get(level, RISK_CFG["Low"])
    c      = cfg["colour"]
    score  = risk["risk_score"]
    conf   = risk.get("confidence", 100)
    warning = risk.get("warning", "")

    # ── Out-of-domain warning banner ──────────────────────────────────────────
    warning_html = ""
    if warning:
        warning_html = f"""
        <div style="background:var(--background-fill-secondary); border-left: 4px solid #ef4444; border-radius:8px;
                    padding:14px 18px; margin-bottom:20px; font-size:0.95em; color:var(--body-text-color); font-weight: 500;">
          {warning}
        </div>"""

    recs_html = "".join(
        f'<li style="margin:8px 0; color:var(--body-text-color);">{r}</li>'
        for r in risk["recommendations"]
    )

    metrics = risk.get("model_metrics", {})
    metrics_html = ""
    if metrics:
        pills = "".join(
            f'<span class="metric-pill" style="font-size: 0.85em; padding: 4px 12px;">{k}: <strong>{v}</strong></span>'
            for k, v in metrics.items() if k != "Model"
        )
        metrics_html = f"""
        <div style="margin-top:24px; padding-top:20px;
                    border-top:1px solid var(--border-color-primary);">
          <div style="font-size:0.8em; text-transform:uppercase; letter-spacing:.1em; font-weight: 700;
                      color:var(--body-text-color-subdued); margin-bottom:12px;">Model Performance</div>
          {pills}
        </div>"""

    return f"""
    <div id="risk-card" style="padding:32px; font-family:'Geist', -apple-system, BlinkMacSystemFont, sans-serif;">
      <div class="section-label" style="margin-top:0;">Risk Assessment</div>

      {warning_html}

      <!-- header -->
      <div style="display:flex; align-items:center; gap:20px; margin-bottom:12px;">
        <div style="display:flex; align-items:center;">
          <span style="display:block; width: 16px; height: 16px; border-radius: 50%; background: {c}; margin-right: 12px;"></span>
          <div>
            <div style="font-size:0.8em; text-transform:uppercase; letter-spacing:.12em;
                        color:{c}; font-weight:800;">Risk Level</div>
            <div style="font-size:2.4em; font-weight:800; color:{c}; line-height:1.2; letter-spacing:-0.02em;">
              {level}
            </div>
          </div>
        </div>
        <div style="margin-left:auto; text-align:right;">
          <div style="font-size:3em; font-weight:900; color:var(--body-text-color); letter-spacing:-0.03em; line-height:1;">{score}</div>
          <div style="font-size:0.9em; color:var(--body-text-color-subdued); font-weight: 600; margin-top:4px;">/ 100</div>
        </div>
      </div>

      {_score_bar(score, c)}

      <!-- confidence row -->
      <div style="display:flex; justify-content:space-between; align-items:center;
                  padding:12px 0 16px; border-bottom:1px solid var(--border-color-primary);">
        <span style="color:var(--body-text-color-subdued); font-size:0.95em; font-weight: 500;">Model Confidence</span>
        <span style="font-weight:700; font-size: 1.15em; color:{'#10b981' if conf >= 50 else '#ef4444'};">
          {conf}%
        </span>
      </div>

      <!-- stats grid -->
      {_stat_row("Flood Coverage",   f"{risk['flood_pct']}%",              c)}
      {_stat_row("Avg Flood Depth",  f"{risk['avg_depth_m']} m",           c)}
      {_stat_row("Max Flood Depth",  f"{depth_info['max_depth_m']} m",     c)}
      {_stat_row("Depth Category",   depth_info['depth_category'],         c)}

      <!-- recommendations -->
      <div style="margin-top:24px; padding-top:20px;
                  border-top:1px solid var(--border-color-primary);">
        <div style="font-size:0.8em; text-transform:uppercase; letter-spacing:.1em; font-weight: 700;
                    color:var(--body-text-color-subdued); margin-bottom:12px;">Recommendations</div>
        <ul style="margin:0; padding-left:24px; line-height:1.8; font-size:0.95em; color:var(--body-text-color);">
          {recs_html}
        </ul>
      </div>

      {metrics_html}
    </div>"""


# ── inference ─────────────────────────────────────────────────────────────────
@spaces.GPU(duration=120)
def predict(image: Image.Image):
    print("=== PREDICT CALLED ===", flush=True)
    if image is None:
        placeholder = f"""
        <div style="background:var(--background-fill-primary); border:2px dashed var(--border-color-primary); border-radius:12px;
                    padding:60px 20px; text-align:center; color:var(--body-text-color-subdued); font-family:'Geist', -apple-system, BlinkMacSystemFont, sans-serif;">
          <div style="display:flex; justify-content:center;"><img src="{logo_src}" style="width: 60px; margin-bottom:16px; opacity: 0.8;"></div>
          <div style="font-size:1.3em; font-weight: 600; color: var(--body-text-color);">Upload an image and click <strong>Analyse</strong></div>
          <div style="font-size:1em; margin-top:10px; color:var(--body-text-color-subdued); font-weight: 400;">
            Supports satellite, aerial, or ground-level flood images
          </div>
        </div>"""
        return None, None, None, None, None, placeholder

    try:
        result     = run_pipeline(image)
        risk       = result["risk"]
        depth_info = result["depth_info"]

        overlay     = Image.fromarray(result["overlay"])
        gradcam     = Image.fromarray(result["gradcam"])
        depth_map   = Image.fromarray(result["depth_map"])
        depth_flood = Image.fromarray(result["depth_overlay"])

        # B&W mask: white = flooded, black = dry — clean and crisp
        mask_bw = Image.fromarray((result["mask"] * 255).astype(np.uint8), mode="L")

        risk_html = build_risk_html(risk, depth_info)
        return overlay, mask_bw, gradcam, depth_map, depth_flood, risk_html
    except Exception as e:
        import traceback
        err_msg = traceback.format_exc()
        print("PREDICTION ERROR:\n" + err_msg)
        error_html = f"<div style='color:var(--body-text-color); background:#fee2e2; border:1px solid #ef4444; border-radius:12px; padding:20px;'><h3 style='color:#b91c1c; margin-top:0;'>Prediction Error</h3><pre style='white-space:pre-wrap; font-size:0.85em;'>{err_msg}</pre></div>"
        return None, None, None, None, None, error_html


# ── UI layout ─────────────────────────────────────────────────────────────────
head_js = """
<script>
    // Always set theme to dark before Gradio initializes
    localStorage.setItem('theme', 'dark');
    document.documentElement.classList.add('dark');
</script>
"""
with gr.Blocks(title="Flood Detection AI") as demo:

    # ── Header ──────────────────────────────────────────────────────────────
    gr.HTML(f"""
    <div id="header-band">
      <label class="theme-switch" title="Toggle Theme">
        <input type="checkbox" onchange="document.documentElement.classList.toggle('dark'); document.documentElement.classList.toggle('light'); document.body.classList.toggle('dark'); document.body.classList.toggle('light'); const c = document.querySelector('.gradio-container'); if(c) {{ c.classList.toggle('dark'); c.classList.toggle('light'); }}" checked>
        <span class="theme-slider"></span>
      </label>
      <h1><img src="{logo_src}" style="width: 40px; vertical-align: middle; margin-right: 12px; margin-bottom: 6px;"> Flood Detection &amp; Risk Assessment</h1>
      <p>
        Attention UNet segmentation &nbsp;·&nbsp; Grad-CAM explainability
        &nbsp;·&nbsp; ZoeDepth metric depth &nbsp;·&nbsp; AI-powered risk scoring
      </p>
      <div style="margin-top:12px;">
        <span class="metric-pill">IoU <strong>76.91%</strong></span>
        <span class="metric-pill">Dice/F1 <strong>86.95%</strong></span>
        <span class="metric-pill">Pixel Acc <strong>89.36%</strong></span>
        <span class="metric-pill">Precision <strong>85.41%</strong></span>
        <span class="metric-pill">Recall <strong>88.54%</strong></span>
      </div>
    </div>
    """)

    # ── Main row: upload + risk card ────────────────────────────────────────
    with gr.Row(equal_height=False):

        with gr.Column(scale=4, elem_id="left-panel"):
            gr.HTML('<div class="section-label" style="margin-top:0;">Input Image</div>')
            input_image = gr.Image(
                type="pil",
                label="",
                elem_classes=["image-card"],
                show_label=False,
            )
            run_btn = gr.Button(
                "Analyse Flood Risk",
                variant="primary",
                elem_id="analyse-btn",
            )
            gr.HTML("""
            <div style="margin-top:20px; padding:20px 24px; background:var(--background-fill-secondary);
                        border-radius:12px; border: 1px solid var(--border-color-primary); border-left:4px solid var(--primary-500);">
              <div style="font-size:0.85em; color:var(--body-text-color); font-weight:800;
                          text-transform:uppercase; letter-spacing:.1em; margin-bottom:12px;">
                Pipeline Steps
              </div>
              <div style="font-size:0.95em; color:var(--body-text-color-subdued); line-height:2.2; font-weight: 500;">
                <span style="opacity: 0.5; margin-right: 8px;">1.</span> Preprocess → 512×512 RGB, normalise<br>
                <span style="opacity: 0.5; margin-right: 8px;">2.</span> Attention UNet → binary flood mask<br>
                <span style="opacity: 0.5; margin-right: 8px;">3.</span> Grad-CAM → attention heatmap<br>
                <span style="opacity: 0.5; margin-right: 8px;">4.</span> ZoeDepth → per-pixel depth (metres)<br>
                <span style="opacity: 0.5; margin-right: 8px;">5.</span> Risk engine → level + score + advice
              </div>
            </div>
            """)

        with gr.Column(scale=6):
            risk_display = gr.HTML(
                value=f"""
                <div style="background:var(--background-fill-primary); border:2px dashed var(--border-color-primary);
                            border-radius:12px; padding:32px; text-align:center;
                            color:var(--body-text-color-subdued); font-family:'Geist', -apple-system, BlinkMacSystemFont, sans-serif; height: 100%; display: flex; flex-direction: column; justify-content: center; min-height: 400px; box-shadow: var(--block-shadow);">
                  <div class="section-label" style="margin-top:0; text-align:left;">Risk Assessment</div>
                  <div style="flex:1; display:flex; flex-direction:column; justify-content:center; align-items:center;">
                    <img src="{logo_src}" style="width: 60px; margin-bottom:16px; opacity: 0.8;">
                    <div style="font-size:1.3em; font-weight: 600; color: var(--body-text-color);">
                      Upload an image and click <strong>Analyse</strong>
                    </div>
                  </div>
                </div>"""
            )

    # ── Visual outputs ───────────────────────────────────────────────────────
    gr.HTML('<div class="section-label">Segmentation &amp; Explainability</div>')
    with gr.Row():
        overlay_out = gr.Image(
            label="Flood Mask Overlay",
            elem_classes=["image-card"],
        )
        mask_bw_out = gr.Image(
            label="Binary Flood Mask (White = Flooded)",
            elem_classes=["image-card"],
            image_mode="L",
        )
        gradcam_out = gr.Image(
            label="Grad-CAM (TURBO)",
            elem_classes=["image-card"],
        )

    gr.HTML('<div class="section-label">ZoeDepth Estimation</div>')
    with gr.Row():
        depth_map_out = gr.Image(
            label="Full Depth Map (PLASMA)",
            elem_classes=["image-card"],
        )
        depth_flood_out = gr.Image(
            label="Flood-Region Depth",
            elem_classes=["image-card"],
        )

    # ── How it works ─────────────────────────────────────────────────────────
    with gr.Accordion("How it works", open=False):
        gr.HTML("""
        <table class="how-table">
          <tr>
            <th>Step</th><th>Component</th><th>What it does</th>
          </tr>
          <tr>
            <td>1</td>
            <td><strong>Preprocessing</strong></td>
            <td>Resize to 512×512, convert to RGB, normalise to [0,1]</td>
          </tr>
          <tr>
            <td>2</td>
            <td><strong>Attention UNet</strong></td>
            <td>Predicts binary flood mask: attention gates suppress irrelevant features</td>
          </tr>
          <tr>
            <td>3</td>
            <td><strong>Grad-CAM</strong></td>
            <td>Gradient-weighted class activation map at conv2d_130: TURBO colormap with contour</td>
          </tr>
          <tr>
            <td>4</td>
            <td><strong>ZoeDepth (Intel/zoedepth-nyu)</strong></td>
            <td>Monocular metric depth estimation: outputs depth in metres per pixel</td>
          </tr>
          <tr>
            <td>5</td>
            <td><strong>Risk Engine</strong></td>
            <td>Combines flood coverage % + avg depth → Low / Moderate / High / Critical</td>
          </tr>
        </table>
        <div style="margin-top:20px; padding:16px 20px; background:var(--background-fill-secondary); border: 1px solid var(--border-color-primary);
                    border-radius:12px; font-size:0.95em; color:var(--body-text-color); font-weight: 500;">
          <strong style="color:var(--body-text-color); margin-right: 12px;">Risk thresholds:</strong>
          <span style="display:inline-block; width:8px; height:8px; border-radius:50%; background:#27ae60; margin-right:6px;"></span> Low (&lt;15% flood, &lt;0.8m) &nbsp;&nbsp;
          <span style="display:inline-block; width:8px; height:8px; border-radius:50%; background:#f39c12; margin-right:6px;"></span> Moderate (15–35%, 0.8–1.5m) &nbsp;&nbsp;
          <span style="display:inline-block; width:8px; height:8px; border-radius:50%; background:#e67e22; margin-right:6px;"></span> High (35–60%, 1.5–2.2m) &nbsp;&nbsp;
          <span style="display:inline-block; width:8px; height:8px; border-radius:50%; background:#e74c3c; margin-right:6px;"></span> Critical (&gt;60% or &gt;2.2m)
        </div>
        """)

    # ── Wire up ───────────────────────────────────────────────────────────────
    run_btn.click(
        fn=predict,
        inputs=[input_image],
        outputs=[overlay_out, mask_bw_out, gradcam_out, depth_map_out, depth_flood_out, risk_display],
    )

if __name__ == "__main__":
    # Pre-download the massive ZoeDepth model to HF cache before starting the server.
    # This prevents the 60-second NGINX timeout when a user makes the first prediction,
    # and avoids the 139 SegFault because we aren't loading PyTorch/TF into memory yet!
    try:
        print("Downloading ZoeDepth weights to cache...")
        from huggingface_hub import snapshot_download
        snapshot_download(repo_id="Intel/zoedepth-nyu")
        print("Download complete. Models are ready!")
    except Exception as e:
        print(f"Warning: Failed to pre-download model: {e}")

    demo.queue().launch(server_name="0.0.0.0", server_port=7860, css=CSS, theme="soft", ssr_mode=False, head=head_js)