Upload app.py

app.py

@@ -0,0 +1,466 @@
+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()
+
+    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
+
+# ── 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)
+
+# ══════════════════════════════════════════════════════════════
+# 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 = 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)
+
+        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)
+
+    if video_file:
+        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!")
+
+        with open(out_path, "rb") as f:
+            st.download_button("⬇️ Download Segmented Video", f,
+                               file_name="geovision_output.mp4", mime="video/mp4")
+
+        if key_frame_results:
+            st.markdown(f"### 🔬 Stage 2 — Key Frame Analysis ({len(key_frame_results)} frames sampled)")
+            for fidx, regions in key_frame_results[:3]:
+                st.markdown(f"**Frame #{fidx}**")
+                show_pipeline_results(regions)
+
+        os.unlink(tmp_path)
+        os.unlink(out_path)
+
+# ══════════════════════════════════════════════════════════════
+# 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 = 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)