app.py

import streamlit as st
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
import cv2
import tempfile
import os
import time
from transformers import (
    SegformerImageProcessor,
    SegformerForSemanticSegmentation,
    AutoImageProcessor,
    SiglipForImageClassification,
)

# ── Page config ───────────────────────────────────────────────
st.set_page_config(
    page_title="GeoVision — Drone Intelligence Platform",
    page_icon="🛰️",
    layout="wide"
)

st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&family=Space+Grotesk:wght@500;700&display=swap');
    html, body, [class*="css"] { font-family: 'Inter', sans-serif; background-color: #0D1117; color: #E6EDF3; }
    .main-title { font-family: 'Space Grotesk', sans-serif; font-size: 2.4rem; font-weight: 700; color: #58A6FF; }
    .sub-title { font-size: 1rem; color: #8B949E; margin-bottom: 1.5rem; }
    .tag { display: inline-block; background: #1F3A5F; color: #58A6FF; border-radius: 20px; padding: 3px 12px;
           font-size: 0.75rem; font-weight: 600; margin-right: 6px; margin-bottom: 16px; }
    .stat-box { background: #161B22; border: 1px solid #30363D; border-radius: 10px; padding: 16px; text-align: center; margin-bottom: 8px; }
    .stat-value { font-family: 'Space Grotesk', sans-serif; font-size: 1.6rem; font-weight: 700; color: #58A6FF; }
    .stat-label { font-size: 0.75rem; color: #8B949E; text-transform: uppercase; letter-spacing: 0.05em; }
    .region-card { background: #161B22; border: 1px solid #30363D; border-radius: 10px; padding: 14px; margin-bottom: 10px; }
    .region-title { font-family: 'Space Grotesk', sans-serif; font-size: 0.95rem; font-weight: 600; color: #E6EDF3; }
    .conf-bar-bg { background: #21262D; border-radius: 4px; height: 6px; margin-top: 6px; }
    .live-badge { display: inline-block; background: #C0392B; color: white; border-radius: 4px;
                  padding: 2px 8px; font-size: 0.7rem; font-weight: 700; }
    .pipeline-badge { display: inline-block; background: #1a472a; color: #2ECC71; border-radius: 4px;
                      padding: 2px 8px; font-size: 0.7rem; font-weight: 700; margin-left: 8px; }
    .footer { text-align: center; color: #484F58; font-size: 0.8rem; margin-top: 3rem;
              padding-top: 1rem; border-top: 1px solid #21262D; }
    .stButton > button { background: linear-gradient(135deg, #1F6FEB, #388BFD); color: white;
                         border: none; border-radius: 8px; padding: 0.6rem 2rem;
                         font-weight: 600; font-size: 1rem; width: 100%; }
</style>
""", unsafe_allow_html=True)

# ── Header ─────────────────────────────────────────────────────
st.markdown('<div class="main-title">🛰️ GeoVision</div>', unsafe_allow_html=True)
st.markdown('<div class="sub-title">Two-Stage Drone Intelligence Pipeline — Land Cover Segmentation + Landform Classification</div>', unsafe_allow_html=True)
st.markdown("""
<span class="tag">🤖 SegFormer-B2</span>
<span class="tag">🔍 SigLIP Landform Classifier</span>
<span class="tag">🌍 Two-Stage AI Pipeline</span>
<span class="tag">📡 Live Feed Ready</span>
""", unsafe_allow_html=True)

# ── Class maps ─────────────────────────────────────────────────
SEG_MAP = {
    0:  ("Structure",    "#7F8C8D"),
    1:  ("Open Area",    "#85C1E9"),
    3:  ("Building",     "#E74C3C"),
    4:  ("Tree",         "#27AE60"),
    6:  ("Road",         "#95A5A6"),
    9:  ("Grass",        "#2ECC71"),
    10: ("Forest",       "#1E8449"),
    12: ("Footpath",     "#D5DBDB"),
    13: ("Terrain",      "#8E44AD"),
    16: ("Bare Land",    "#D4AC0D"),
    17: ("Water",        "#3498DB"),
    21: ("Farmland",     "#F39C12"),
    29: ("Open Field",   "#F0B27A"),
    43: ("Sand/Desert",  "#F7DC6F"),
    46: ("Barren",       "#CA6F1E"),
    94: ("Vegetation",   "#1ABC9C"),
}

LANDFORM_LABELS = {
    "0": "Annual Crop",
    "1": "Forest",
    "2": "Herbaceous Vegetation",
    "3": "Highway",
    "4": "Industrial",
    "5": "Pasture",
    "6": "Permanent Crop",
    "7": "Residential",
    "8": "River",
    "9": "Sea / Lake"
}

LANDFORM_ICONS = {
    "Annual Crop": "🌾",
    "Forest": "🌲",
    "Herbaceous Vegetation": "🌿",
    "Highway": "🛣️",
    "Industrial": "🏭",
    "Pasture": "🐄",
    "Permanent Crop": "🍇",
    "Residential": "🏘️",
    "River": "🏞️",
    "Sea / Lake": "🌊"
}

GEOLOGICAL_HINTS = {
    "Annual Crop":           "Likely alluvial/loam soil. Good agricultural fertility.",
    "Forest":                "Humid soil with organic layer. Possible clay-rich substrate.",
    "Herbaceous Vegetation": "Thin topsoil over sedimentary or volcanic base.",
    "Highway":               "Compacted sub-base. Engineered ground — no natural geology.",
    "Industrial":            "Artificially modified land. Possible fill or compacted gravel.",
    "Pasture":               "Silty or clay-loam soil. Generally flat sedimentary terrain.",
    "Permanent Crop":        "Well-drained loamy soil. Likely sedimentary or alluvial.",
    "Residential":           "Urban fill. Natural geology obscured.",
    "River":                 "Active alluvial deposit. Sandy/gravelly riverbed sediment.",
    "Sea / Lake":            "Water body. Lakebed may have clay, silt, or sand deposits.",
}

def hex_to_rgb(h):
    h = h.lstrip("#")
    return tuple(int(h[i:i+2], 16) for i in (0, 2, 4))

def cid_to_rgb(cid):
    return hex_to_rgb(SEG_MAP[cid][1]) if cid in SEG_MAP else (189, 195, 199)

# ── Model loaders ──────────────────────────────────────────────
@st.cache_resource(show_spinner=False)
def load_seg_model():
    proc = SegformerImageProcessor.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512")
    mdl  = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512")
    dev  = "cuda" if torch.cuda.is_available() else "cpu"
    return proc, mdl.to(dev).eval(), dev

@st.cache_resource(show_spinner=False)
def load_cls_model():
    proc = AutoImageProcessor.from_pretrained("prithivMLmods/SAT-Landforms-Classifier")
    mdl  = SiglipForImageClassification.from_pretrained("prithivMLmods/SAT-Landforms-Classifier")
    dev  = "cuda" if torch.cuda.is_available() else "cpu"
    return proc, mdl.to(dev).eval(), dev

# ── Stage 1: Segmentation ──────────────────────────────────────
def run_segmentation(image: Image.Image):
    proc, mdl, dev = load_seg_model()
    img = image.resize((512, 512))
    inputs = proc(images=img, return_tensors="pt").to(dev)
    with torch.no_grad():
        logits = mdl(**inputs).logits

    up   = torch.nn.functional.interpolate(logits, size=(512, 512), mode="bilinear", align_corners=False)
    pred = up.argmax(dim=1)[0].cpu().numpy()

    # Confidence heatmap — uncertainty = 1 - max softmax prob per pixel
    probs       = torch.nn.functional.softmax(up, dim=1)
    max_probs   = probs.max(dim=1)[0][0].cpu().numpy()
    uncertainty = 1.0 - max_probs  # high = uncertain, low = confident

    color_map = np.zeros((512, 512, 3), dtype=np.uint8)
    for cid in np.unique(pred):
        color_map[pred == cid] = cid_to_rgb(cid)

    seg_img  = Image.fromarray(color_map)
    blended  = Image.blend(img, seg_img, alpha=0.55)

    total = 512 * 512
    stats = {}
    for cid in np.unique(pred):
        pct  = round(np.sum(pred == cid) / total * 100, 1)
        name = SEG_MAP.get(cid, ("Other",))[0]
        stats[name] = stats.get(name, 0) + pct

    return img, seg_img, blended, stats, pred, uncertainty

# ── Stage 2: Per-region landform classification ────────────────
def classify_regions(orig_img: Image.Image, pred: np.ndarray):
    proc, mdl, dev = load_cls_model()
    orig_arr = np.array(orig_img)
    results  = []

    detected_cids = [cid for cid in np.unique(pred) if cid in SEG_MAP]

    for cid in detected_cids:
        mask = pred == cid
        ys, xs = np.where(mask)
        if len(ys) < 500:   # skip tiny regions
            continue

        # Crop bounding box of this region
        y1, y2 = ys.min(), ys.max()
        x1, x2 = xs.min(), xs.max()
        patch = orig_arr[y1:y2+1, x1:x2+1]
        if patch.size == 0:
            continue
        patch_pil = Image.fromarray(patch).resize((224, 224))

        inputs = proc(images=patch_pil, return_tensors="pt").to(dev)
        with torch.no_grad():
            logits = mdl(**inputs).logits
        probs  = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
        top_id = int(np.argmax(probs))
        top_lf = LANDFORM_LABELS[str(top_id)]
        top_conf = round(probs[top_id] * 100, 1)

        seg_name  = SEG_MAP[cid][0]
        seg_color = SEG_MAP[cid][1]
        icon      = LANDFORM_ICONS.get(top_lf, "📍")
        geo_hint  = GEOLOGICAL_HINTS.get(top_lf, "")

        results.append({
            "seg_name":  seg_name,
            "seg_color": seg_color,
            "landform":  top_lf,
            "icon":      icon,
            "conf":      top_conf,
            "geo_hint":  geo_hint,
            "all_probs": {LANDFORM_LABELS[str(i)]: round(probs[i]*100,1) for i in range(len(probs))},
        })

    results.sort(key=lambda x: -x["conf"])
    return results

# ── Display helpers ────────────────────────────────────────────
def show_images(orig, seg, blend):
    c1, c2, c3 = st.columns(3)
    with c1: st.image(orig,  caption="📷 Original Drone Image",  use_column_width=True)
    with c2: st.image(seg,   caption="🗺️ Segmentation Map",      use_column_width=True)
    with c3: st.image(blend, caption="🔀 Blended Overlay",       use_column_width=True)

def show_stats(stats):
    top = sorted(stats.items(), key=lambda x: -x[1])[:4]
    cols = st.columns(4)
    for i, (name, pct) in enumerate(top):
        with cols[i]:
            st.markdown(f"""
            <div class="stat-box">
                <div class="stat-value">{pct}%</div>
                <div class="stat-label">{name}</div>
            </div>""", unsafe_allow_html=True)

def show_pipeline_results(regions):
    st.markdown("---")
    st.markdown("""
    ### 🔬 Stage 2 — Landform & Geological Intelligence
    <span class="pipeline-badge">✦ TWO-STAGE AI PIPELINE</span>
    <p style='color:#8B949E; font-size:0.9rem; margin-top:8px;'>
    Each land cover region detected by SegFormer is independently classified by a satellite landform classifier.
    Geological context is inferred from the landform type.
    </p>
    """, unsafe_allow_html=True)

    if not regions:
        st.info("No significant regions detected for Stage 2 classification.")
        return

    cols = st.columns(2)
    for i, r in enumerate(regions):
        with cols[i % 2]:
            conf_color = "#2ECC71" if r["conf"] > 60 else "#F39C12" if r["conf"] > 35 else "#E74C3C"
            st.markdown(f"""
            <div class="region-card">
                <div style="display:flex; align-items:center; gap:10px; margin-bottom:8px;">
                    <div style="width:14px;height:14px;border-radius:3px;background:{r['seg_color']};flex-shrink:0;"></div>
                    <span class="region-title">{r['seg_name']}</span>
                    <span style="margin-left:auto;font-size:0.75rem;color:#8B949E;">{r['conf']}% conf.</span>
                </div>
                <div style="font-size:1.4rem;">{r['icon']} <span style="font-size:0.95rem;font-weight:600;color:#E6EDF3;">{r['landform']}</span></div>
                <div class="conf-bar-bg">
                    <div style="width:{min(r['conf'],100)}%;background:{conf_color};height:6px;border-radius:4px;"></div>
                </div>
                <div style="font-size:0.8rem;color:#8B949E;margin-top:8px;">🪨 <i>{r['geo_hint']}</i></div>
            </div>
            """, unsafe_allow_html=True)

def show_summary_chart(regions):
    if not regions:
        return
    st.markdown("### 📊 Landform Distribution")
    lf_counts = {}
    for r in regions:
        lf_counts[r["landform"]] = lf_counts.get(r["landform"], 0) + 1

    fig, ax = plt.subplots(figsize=(10, 3))
    fig.patch.set_facecolor("#0D1117")
    ax.set_facecolor("#161B22")
    names = list(lf_counts.keys())
    vals  = list(lf_counts.values())
    colors = ["#58A6FF","#2ECC71","#E74C3C","#F39C12","#8E44AD","#3498DB"]
    ax.barh(names, vals, color=colors[:len(names)], height=0.5, edgecolor="none")
    ax.set_xlabel("Region Count", color="#8B949E")
    ax.tick_params(colors="#8B949E")
    ax.spines[:].set_color("#30363D")
    plt.tight_layout()
    st.pyplot(fig)

def show_csv_export(stats):
    """Generate and offer CSV download of land cover stats."""
    import io
    st.markdown("### 📥 Export Results")
    lines = ["Land Cover Class,Coverage (%),Coverage Category"]
    for name, pct in sorted(stats.items(), key=lambda x: -x[1]):
        if pct >= 10:
            category = "Dominant"
        elif pct >= 5:
            category = "Significant"
        else:
            category = "Minor"
        lines.append(f"{name},{pct},{category}")
    csv_str = "\n".join(lines)
    st.download_button(
        label="⬇️ Download Land Cover Report (CSV)",
        data=csv_str,
        file_name="geovision_land_cover_report.csv",
        mime="text/csv",
    )

def show_uncertainty_heatmap(uncertainty: np.ndarray, orig: Image.Image):
    """Show confidence heatmap — red = uncertain, blue = confident."""
    st.markdown("### 🌡️ Model Confidence Heatmap")
    st.caption("Red areas = model is uncertain about classification · Blue/dark = high confidence")

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    fig.patch.set_facecolor("#0D1117")

    # Original
    axes[0].imshow(orig)
    axes[0].set_title("Original Image", color="#E6EDF3", fontsize=12)
    axes[0].axis("off")

    # Uncertainty heatmap
    im = axes[1].imshow(uncertainty, cmap="RdYlBu_r", vmin=0, vmax=1)
    axes[1].set_title("Uncertainty Map", color="#E6EDF3", fontsize=12)
    axes[1].axis("off")

    cbar = fig.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04)
    cbar.set_label("Uncertainty (0=confident, 1=uncertain)", color="#8B949E", fontsize=9)
    cbar.ax.yaxis.set_tick_params(color="#8B949E")
    plt.setp(cbar.ax.yaxis.get_ticklabels(), color="#8B949E")

    # Overall confidence score
    avg_conf = round((1 - uncertainty.mean()) * 100, 1)
    axes[1].set_xlabel(f"Overall Model Confidence: {avg_conf}%",
                       color="#58A6FF", fontsize=11, fontweight="bold")

    fig.patch.set_facecolor("#0D1117")
    plt.tight_layout()
    st.pyplot(fig)
    st.markdown(f"""
    <div class="stat-box" style="margin-top:8px;">
        <div class="stat-value">{avg_conf}%</div>
        <div class="stat-label">Overall Segmentation Confidence</div>
    </div>
    """, unsafe_allow_html=True)

# ══════════════════════════════════════════════════════════════
# MODE SELECTOR
# ══════════════════════════════════════════════════════════════
st.markdown("---")
mode = st.radio("**Select Input Mode:**",
    ["📷 Image", "🎥 Video", "📡 Live Webcam Feed"], horizontal=True)
st.markdown("---")

# ══════════════════════════════════════════════════════════════
# MODE 1: IMAGE
# ══════════════════════════════════════════════════════════════
if mode == "📷 Image":
    st.markdown("#### 📷 Upload Drone / Aerial Image")
    uploaded = st.file_uploader("", type=["jpg","jpeg","png","tif","tiff"])

    if uploaded:
        image = Image.open(uploaded).convert("RGB")

        with st.spinner("🔍 Stage 1 — Running land cover segmentation..."):
            orig, seg, blend, stats, pred, uncertainty = run_segmentation(image)

        st.markdown("### 🗺️ Stage 1 — Land Cover Segmentation")
        show_stats(stats)
        st.markdown("<br>", unsafe_allow_html=True)
        show_images(orig, seg, blend)

        # CSV Export
        show_csv_export(stats)

        # Uncertainty heatmap
        show_uncertainty_heatmap(uncertainty, orig)

        with st.spinner("🔬 Stage 2 — Classifying each region's landform & geology..."):
            regions = classify_regions(orig, pred)

        show_pipeline_results(regions)
        show_summary_chart(regions)

# ══════════════════════════════════════════════════════════════
# MODE 2: VIDEO
# ══════════════════════════════════════════════════════════════
elif mode == "🎥 Video":
    st.markdown("#### 🎥 Upload Drone Video")
    st.info("Uploads a drone video and segments every Nth frame. Stage 2 classification runs on key frames.")

    video_file  = st.file_uploader("", type=["mp4","avi","mov","mkv"])
    frame_skip  = st.slider("Process every N frames", 5, 60, 15)
    run_stage2  = st.checkbox("Also run Stage 2 landform classification on key frames", value=True)

    # Init session state for video
    if "video_done" not in st.session_state:
        st.session_state.video_done = False
    if "video_bytes" not in st.session_state:
        st.session_state.video_bytes = None
    if "video_key_frames" not in st.session_state:
        st.session_state.video_key_frames = []
    if "last_video_name" not in st.session_state:
        st.session_state.last_video_name = None

    # Only process if new video uploaded
    if video_file:
        if video_file.name != st.session_state.last_video_name:
            # Reset state for new video
            st.session_state.video_done = False
            st.session_state.video_bytes = None
            st.session_state.video_key_frames = []
            st.session_state.last_video_name = video_file.name

        if not st.session_state.video_done:
            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp:
                tmp.write(video_file.read())
                tmp_path = tmp.name

            cap          = cv2.VideoCapture(tmp_path)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            fps          = cap.get(cv2.CAP_PROP_FPS) or 24
            out_path     = tmp_path.replace(".mp4", "_geo.mp4")
            fourcc       = cv2.VideoWriter_fourcc(*"mp4v")
            out          = cv2.VideoWriter(out_path, fourcc, max(fps/frame_skip, 1), (512, 512))

            progress  = st.progress(0, text="Processing...")
            frame_idx = 0
            key_frame_results = []

            proc_seg, mdl_seg, dev_seg = load_seg_model()

            with st.spinner("Processing video..."):
                while True:
                    ret, frame = cap.read()
                    if not ret:
                        break
                    if frame_idx % frame_skip == 0:
                        rgb    = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                        pil    = Image.fromarray(rgb).resize((512, 512))
                        inputs = proc_seg(images=pil, return_tensors="pt").to(dev_seg)
                        with torch.no_grad():
                            logits = mdl_seg(**inputs).logits
                        up   = torch.nn.functional.interpolate(logits, size=(512,512), mode="bilinear", align_corners=False)
                        pred = up.argmax(dim=1)[0].cpu().numpy()

                        color_map = np.zeros((512,512,3), dtype=np.uint8)
                        for cid in np.unique(pred):
                            color_map[pred==cid] = cid_to_rgb(cid)

                        color_bgr    = cv2.cvtColor(color_map, cv2.COLOR_RGB2BGR)
                        frame_resize = cv2.resize(frame, (512,512))
                        blended_cv   = cv2.addWeighted(frame_resize, 0.5, color_bgr, 0.5, 0)
                        out.write(blended_cv)

                        if run_stage2 and frame_idx % (frame_skip * 5) == 0:
                            regions = classify_regions(pil, pred)
                            if regions:
                                key_frame_results.append((frame_idx, regions))

                        progress.progress(min(frame_idx/max(total_frames,1), 1.0),
                                          text=f"Frame {frame_idx}/{total_frames}")
                    frame_idx += 1

            cap.release()
            out.release()
            progress.progress(1.0, text="✅ Done!")

            # Save results to session state
            with open(out_path, "rb") as f:
                st.session_state.video_bytes = f.read()
            st.session_state.video_key_frames = key_frame_results
            st.session_state.video_done = True

            os.unlink(tmp_path)
            os.unlink(out_path)

        # Always show results from session state — no reprocessing
        if st.session_state.video_done and st.session_state.video_bytes:
            st.success("✅ Video processed successfully!")
            st.download_button(
                "⬇️ Download Segmented Video",
                st.session_state.video_bytes,
                file_name="geovision_output.mp4",
                mime="video/mp4"
            )

            if st.session_state.video_key_frames:
                st.markdown(f"### 🔬 Stage 2 — Key Frame Analysis ({len(st.session_state.video_key_frames)} frames sampled)")
                for fidx, regions in st.session_state.video_key_frames[:3]:
                    st.markdown(f"**Frame #{fidx}**")
                    show_pipeline_results(regions)

# ══════════════════════════════════════════════════════════════
# MODE 3: LIVE WEBCAM
# ══════════════════════════════════════════════════════════════
elif mode == "📡 Live Webcam Feed":
    st.markdown(
        '#### 📡 Live Drone Feed &nbsp;<span class="live-badge">● LIVE</span>'
        '&nbsp;<span class="pipeline-badge">✦ TWO-STAGE</span>',
        unsafe_allow_html=True
    )
    st.info("Simulates a live drone feed. Each frame is processed through the full two-stage pipeline.")

    col1, col2 = st.columns([1,2])
    with col1:
        interval   = st.slider("Capture interval (seconds)", 2, 15, 5)
        run_stage2 = st.checkbox("Enable Stage 2 classification", value=True)
        run_live   = st.checkbox("▶️ Start Live Feed", value=False)
    with col2:
        st.markdown("""
        **Two-stage pipeline on every frame:**
        1. SegFormer segments the frame into land cover regions
        2. Each region is classified for landform type
        3. Geological context is inferred in real time
        """)

    if run_live:
        cap = cv2.VideoCapture(0)
        if not cap.isOpened():
            st.error("❌ Could not access webcam. Please allow camera permissions.")
        else:
            frame_ph  = st.empty()
            img_ph    = st.empty()
            region_ph = st.empty()
            count     = 0

            while run_live:
                ret, frame = cap.read()
                if not ret:
                    break
                count += 1

                rgb  = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                pil  = Image.fromarray(rgb).resize((512, 512))
                orig, seg, blend, stats, pred, uncertainty = run_segmentation(pil)

                with img_ph.container():
                    frame_ph.markdown(f"**📡 Live Frame #{count}**")
                    c1, c2, c3 = st.columns(3)
                    with c1: st.image(orig,  caption="Original",    use_column_width=True)
                    with c2: st.image(seg,   caption="Segmentation", use_column_width=True)
                    with c3: st.image(blend, caption="Overlay",     use_column_width=True)
                    show_stats(stats)

                if run_stage2:
                    regions = classify_regions(orig, pred)
                    with region_ph.container():
                        show_pipeline_results(regions)

                time.sleep(interval)

            cap.release()

# ── Footer ─────────────────────────────────────────────────────
st.markdown("""
<div class="footer">
    GeoVision · Two-Stage Geospatial AI Pipeline<br>
    SegFormer-B2 (Land Cover) + SigLIP (Landform Classification) · Built with 🤗 Transformers & Streamlit
</div>
""", unsafe_allow_html=True)