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("""
""", unsafe_allow_html=True)
# ── Header ─────────────────────────────────────────────────────
st.markdown('🛰️ GeoVision
', unsafe_allow_html=True)
st.markdown('Two-Stage Drone Intelligence Pipeline — Land Cover Segmentation + Landform Classification
', unsafe_allow_html=True)
st.markdown("""
🤖 SegFormer-B2
🔍 SigLIP Landform Classifier
🌍 Two-Stage AI Pipeline
📡 Live Feed Ready
""", 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"""
""", unsafe_allow_html=True)
def show_pipeline_results(regions):
st.markdown("---")
st.markdown("""
### 🔬 Stage 2 — Landform & Geological Intelligence
✦ TWO-STAGE AI PIPELINE
Each land cover region detected by SegFormer is independently classified by a satellite landform classifier.
Geological context is inferred from the landform type.
""", 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"""
{r['seg_name']}
{r['conf']}% conf.
{r['icon']} {r['landform']}
🪨 {r['geo_hint']}
""", 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"""
{avg_conf}%
Overall Segmentation Confidence
""", 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("
", 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 ● LIVE'
' ✦ TWO-STAGE',
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("""
""", unsafe_allow_html=True)