Spaces:
Sleeping
Sleeping
| """Janus — Feature Extraction via SAM 3 with Tiling + Live Dashboard. | |
| Upload large GeoTIFFs, select feature types, and watch as SAM 3 processes | |
| tile by tile with live progress and confidence tracking. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| import os | |
| import tempfile | |
| import time | |
| from collections import defaultdict | |
| from pathlib import Path | |
| import gradio as gr | |
| import geopandas as gpd | |
| import spaces | |
| import matplotlib | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import rasterio | |
| import torch | |
| from affine import Affine | |
| from PIL import Image | |
| from pyproj import CRS | |
| from rasterio.features import shapes | |
| from rasterio.plot import show | |
| from rasterio.windows import Window | |
| from shapely.geometry import shape | |
| # --------------------------------------------------------------------------- | |
| # Globals | |
| # --------------------------------------------------------------------------- | |
| MODEL = None | |
| PROCESSOR = None | |
| def get_device(): | |
| """Detect GPU at runtime, not import time.""" | |
| import subprocess | |
| print(f"[GPU DEBUG] torch.cuda.is_available() = {torch.cuda.is_available()}") | |
| print(f"[GPU DEBUG] torch.version.cuda = {torch.version.cuda}") | |
| print(f"[GPU DEBUG] torch.backends.cudnn.enabled = {torch.backends.cudnn.enabled}") | |
| try: | |
| print(f"[GPU DEBUG] torch.cuda.device_count() = {torch.cuda.device_count()}") | |
| except Exception as e: | |
| print(f"[GPU DEBUG] torch.cuda.device_count() ERROR: {e}") | |
| try: | |
| result = subprocess.run(["nvidia-smi"], capture_output=True, text=True, timeout=5) | |
| print(f"[GPU DEBUG] nvidia-smi output:\n{result.stdout[:500]}") | |
| if result.stderr: | |
| print(f"[GPU DEBUG] nvidia-smi stderr: {result.stderr[:200]}") | |
| except Exception as e: | |
| print(f"[GPU DEBUG] nvidia-smi ERROR: {e}") | |
| try: | |
| import os | |
| print(f"[GPU DEBUG] CUDA_VISIBLE_DEVICES = {os.environ.get('CUDA_VISIBLE_DEVICES', 'NOT SET')}") | |
| print(f"[GPU DEBUG] NVIDIA_VISIBLE_DEVICES = {os.environ.get('NVIDIA_VISIBLE_DEVICES', 'NOT SET')}") | |
| except Exception as e: | |
| print(f"[GPU DEBUG] env check ERROR: {e}") | |
| if torch.cuda.is_available(): | |
| name = torch.cuda.get_device_name(0) | |
| print(f"[GPU DEBUG] Using GPU: {name}") | |
| return "cuda" | |
| print("[GPU DEBUG] Falling back to CPU") | |
| return "cpu" | |
| matplotlib.use("Agg") | |
| # --------------------------------------------------------------------------- | |
| # Feature presets | |
| # --------------------------------------------------------------------------- | |
| FEATURE_PRESETS = { | |
| "Building": { | |
| "prompt": "building", | |
| "min_area": 20.0, "max_area": 50000.0, | |
| "min_compactness": 0.25, "min_rectangularity": 0.5, | |
| "color": "#10b981", | |
| }, | |
| "Road": { | |
| "prompt": "road", | |
| "min_area": 10.0, "max_area": 500000.0, | |
| "min_compactness": 0.0, "min_rectangularity": 0.0, | |
| "color": "#f59e0b", | |
| }, | |
| "Waterbody": { | |
| "prompt": "water", | |
| "min_area": 50.0, "max_area": 5000000.0, | |
| "min_compactness": 0.0, "min_rectangularity": 0.0, | |
| "color": "#3b82f6", | |
| }, | |
| "Vegetation": { | |
| "prompt": "tree", | |
| "min_area": 30.0, "max_area": 5000000.0, | |
| "min_compactness": 0.0, "min_rectangularity": 0.0, | |
| "color": "#22c55e", | |
| }, | |
| "Parking Lot": { | |
| "prompt": "parking lot", | |
| "min_area": 100.0, "max_area": 200000.0, | |
| "min_compactness": 0.2, "min_rectangularity": 0.4, | |
| "color": "#8b5cf6", | |
| }, | |
| } | |
| WORLD_EXT_MAP = { | |
| ".png": ".pgw", ".jpg": ".jgw", ".jpeg": ".jgw", | |
| ".tif": ".tfw", ".tiff": ".tfw", | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Model | |
| # --------------------------------------------------------------------------- | |
| def load_model(): | |
| global MODEL, PROCESSOR | |
| if MODEL is None: | |
| from transformers import Sam3Model, Sam3Processor | |
| device = get_device() | |
| print(f"[MODEL] Loading SAM 3 on {device}...") | |
| MODEL = Sam3Model.from_pretrained("facebook/sam3").to(device) | |
| PROCESSOR = Sam3Processor.from_pretrained("facebook/sam3") | |
| print(f"[MODEL] SAM 3 loaded on {device}") | |
| return MODEL, PROCESSOR | |
| # --------------------------------------------------------------------------- | |
| # Ingest | |
| # --------------------------------------------------------------------------- | |
| def _parse_world_file(world_path: str) -> Affine: | |
| lines = Path(world_path).read_text().strip().splitlines() | |
| if len(lines) < 6: | |
| raise gr.Error(f"World file must have 6 lines, got {len(lines)}") | |
| return Affine( | |
| float(lines[0]), float(lines[2]), float(lines[4]), | |
| float(lines[1]), float(lines[3]), float(lines[5]), | |
| ) | |
| def ingest(image_path: str, world_path: str | None, crs: str) -> str: | |
| tmp = tempfile.mkdtemp(prefix="janus_") | |
| geotiff = os.path.join(tmp, "input.tif") | |
| if image_path.lower().endswith((".tif", ".tiff")): | |
| try: | |
| with rasterio.open(image_path) as src: | |
| if src.crs is not None and src.transform != Affine.identity(): | |
| import shutil | |
| shutil.copy2(image_path, geotiff) | |
| return geotiff | |
| except Exception: | |
| pass | |
| if world_path is None: | |
| img_p = Path(image_path) | |
| expected_ext = WORLD_EXT_MAP.get(img_p.suffix.lower(), ".pgw") | |
| auto_wf = img_p.with_suffix(expected_ext) | |
| if auto_wf.exists(): | |
| world_path = str(auto_wf) | |
| else: | |
| raise gr.Error(f"No world file found. Expected: {auto_wf.name}") | |
| transform = _parse_world_file(world_path) | |
| target_crs = CRS.from_user_input(crs) | |
| img = Image.open(image_path).convert("RGB") | |
| img_array = np.array(img) | |
| h, w, bands = img_array.shape | |
| profile = { | |
| "driver": "GTiff", "dtype": "uint8", | |
| "width": w, "height": h, "count": bands, | |
| "crs": target_crs, "transform": transform, | |
| } | |
| with rasterio.open(geotiff, "w", **profile) as dst: | |
| for b in range(bands): | |
| dst.write(img_array[:, :, b], b + 1) | |
| return geotiff | |
| # --------------------------------------------------------------------------- | |
| # Tiling | |
| # --------------------------------------------------------------------------- | |
| def compute_tile_windows(img_w: int, img_h: int, tile_size: int, overlap: int) -> list[Window]: | |
| step = tile_size - overlap | |
| windows = [] | |
| for y in range(0, img_h, step): | |
| for x in range(0, img_w, step): | |
| tw = min(tile_size, img_w - x) | |
| th = min(tile_size, img_h - y) | |
| if tw < tile_size // 4 or th < tile_size // 4: | |
| continue | |
| windows.append(Window(x, y, tw, th)) | |
| return windows | |
| # --------------------------------------------------------------------------- | |
| # Per-tile segmentation | |
| # --------------------------------------------------------------------------- | |
| def segment_tile(tile_rgb: np.ndarray, prompt: str, confidence: float): | |
| model, processor = load_model() | |
| device = get_device() | |
| image = Image.fromarray(tile_rgb) | |
| inputs = processor(images=image, text=prompt, return_tensors="pt").to(device) | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| results = processor.post_process_instance_segmentation( | |
| outputs, threshold=confidence, mask_threshold=0.5, | |
| target_sizes=inputs.get("original_sizes").tolist(), | |
| )[0] | |
| return results["masks"].cpu().numpy(), results["scores"].cpu().numpy() | |
| # --------------------------------------------------------------------------- | |
| # Shape metrics | |
| # --------------------------------------------------------------------------- | |
| def _compactness(geom): | |
| if geom.is_empty or geom.length == 0: | |
| return 0.0 | |
| return (4.0 * math.pi * geom.area) / (geom.length ** 2) | |
| def _rectangularity(geom): | |
| if geom.is_empty: | |
| return 0.0 | |
| mrr = geom.minimum_rotated_rectangle | |
| return geom.area / mrr.area if mrr.area > 0 else 0.0 | |
| # --------------------------------------------------------------------------- | |
| # Dedup + filter | |
| # --------------------------------------------------------------------------- | |
| def merge_and_deduplicate(features: list[dict], crs_obj, iou_thresh: float = 0.5) -> gpd.GeoDataFrame: | |
| if not features: | |
| return gpd.GeoDataFrame(columns=["feature_id", "feature_type", "confidence", "geometry"]) | |
| gdf = gpd.GeoDataFrame(features, crs=crs_obj) | |
| if len(gdf) == 0: | |
| return gdf | |
| sindex = gdf.sindex | |
| drop = set() | |
| for idx, row in gdf.iterrows(): | |
| if idx in drop: | |
| continue | |
| for cand_idx in sindex.intersection(row.geometry.bounds): | |
| if cand_idx <= idx or cand_idx in drop: | |
| continue | |
| cand = gdf.loc[cand_idx] | |
| if row["feature_type"] != cand["feature_type"]: | |
| continue | |
| if not row.geometry.intersects(cand.geometry): | |
| continue | |
| try: | |
| inter = row.geometry.intersection(cand.geometry).area | |
| union = row.geometry.union(cand.geometry).area | |
| if union > 0 and inter / union >= iou_thresh: | |
| if row["confidence"] >= cand["confidence"]: | |
| drop.add(cand_idx) | |
| else: | |
| drop.add(idx) | |
| break | |
| except Exception: | |
| continue | |
| gdf = gdf.drop(index=drop).reset_index(drop=True) | |
| gdf["feature_id"] = range(1, len(gdf) + 1) | |
| return gdf | |
| def filter_features(gdf: gpd.GeoDataFrame, preset: dict) -> gpd.GeoDataFrame: | |
| if len(gdf) == 0: | |
| return gdf | |
| if gdf.crs and gdf.crs.is_geographic: | |
| proj = gdf.to_crs(gdf.estimate_utm_crs()) | |
| gdf["area_m2"] = proj.geometry.area | |
| pg = proj.geometry | |
| else: | |
| gdf["area_m2"] = gdf.geometry.area | |
| pg = gdf.geometry | |
| gdf = gdf[(gdf["area_m2"] >= preset["min_area"]) & (gdf["area_m2"] <= preset["max_area"])].copy() | |
| if len(gdf) == 0: | |
| return gdf | |
| if gdf.crs and gdf.crs.is_geographic: | |
| pg = gdf.to_crs(gdf.estimate_utm_crs()).geometry | |
| else: | |
| pg = gdf.geometry | |
| gdf["compactness"] = pg.apply(_compactness) | |
| if preset["min_compactness"] > 0: | |
| gdf = gdf[gdf["compactness"] >= preset["min_compactness"]].copy() | |
| if len(gdf) == 0: | |
| return gdf | |
| if preset["min_rectangularity"] > 0: | |
| if gdf.crs and gdf.crs.is_geographic: | |
| pg = gdf.to_crs(gdf.estimate_utm_crs()).geometry | |
| else: | |
| pg = gdf.geometry | |
| gdf["rectangularity"] = pg.apply(_rectangularity) | |
| gdf = gdf[gdf["rectangularity"] >= preset["min_rectangularity"]].copy() | |
| return gdf | |
| # --------------------------------------------------------------------------- | |
| # Export | |
| # --------------------------------------------------------------------------- | |
| def export_all(features_by_type: dict[str, gpd.GeoDataFrame], tmp_dir: str) -> list[str]: | |
| paths = [] | |
| for feat_type, gdf in features_by_type.items(): | |
| if len(gdf) == 0: | |
| continue | |
| name = feat_type.lower().replace(" ", "_") | |
| gpkg = os.path.join(tmp_dir, f"{name}.gpkg") | |
| gdf.to_file(gpkg, driver="GPKG") | |
| paths.append(gpkg) | |
| gj = os.path.join(tmp_dir, f"{name}.geojson") | |
| gdf.to_file(gj, driver="GeoJSON") | |
| paths.append(gj) | |
| wkt = os.path.join(tmp_dir, f"{name}.wkt") | |
| with open(wkt, "w") as f: | |
| for idx, row in enumerate(gdf.itertuples(), start=1): | |
| c = getattr(row, "confidence", 0.0) | |
| f.write(f"{idx}|{c:.3f}|{row.geometry.wkt}\n") | |
| paths.append(wkt) | |
| return paths | |
| # --------------------------------------------------------------------------- | |
| # Overlay | |
| # --------------------------------------------------------------------------- | |
| def make_overlay( | |
| geotiff: str, | |
| features_by_type: dict[str, gpd.GeoDataFrame], | |
| scanned_bounds: tuple | None = None, | |
| ) -> str: | |
| """Render before/after overlay with optional scan-progress box. | |
| Args: | |
| geotiff: Path to the source GeoTIFF. | |
| features_by_type: Dict of feature_type -> GeoDataFrame. | |
| scanned_bounds: (left, bottom, right, top) in CRS coordinates showing | |
| the area processed so far. Drawn as a dashed box. | |
| """ | |
| fig, axes = plt.subplots(1, 2, figsize=(16, 7), dpi=100) | |
| fig.patch.set_facecolor("#fafafa") | |
| with rasterio.open(geotiff) as src: | |
| max_dim = 2000 | |
| scale = min(1.0, max_dim / max(src.width, src.height)) | |
| out_w = int(src.width * scale) | |
| out_h = int(src.height * scale) | |
| data = src.read( | |
| out_shape=(src.count, out_h, out_w), | |
| resampling=rasterio.enums.Resampling.bilinear, | |
| ) | |
| extent = [src.bounds.left, src.bounds.right, src.bounds.bottom, src.bounds.top] | |
| rgb = np.transpose(data[:3], (1, 2, 0)) | |
| # Left panel: input image | |
| axes[0].imshow(rgb, extent=extent) | |
| axes[0].set_title("Input", fontsize=14, fontweight=600, color="#18181b", pad=12) | |
| axes[0].tick_params(labelsize=7, colors="#71717a") | |
| # Right panel: image + detected features | |
| axes[1].imshow(rgb, extent=extent) | |
| # Draw scan progress box | |
| if scanned_bounds is not None: | |
| from matplotlib.patches import Rectangle | |
| left, bottom, right, top = scanned_bounds | |
| rect = Rectangle( | |
| (left, bottom), right - left, top - bottom, | |
| linewidth=1.5, edgecolor="#71717a", facecolor="none", | |
| linestyle="--", alpha=0.6, | |
| ) | |
| axes[1].add_patch(rect) | |
| # Plot features color-coded by type | |
| total = 0 | |
| legend_items = [] | |
| for feat_type, gdf in features_by_type.items(): | |
| if len(gdf) == 0: | |
| continue | |
| color = FEATURE_PRESETS.get(feat_type, {}).get("color", "#ef4444") | |
| gdf.plot(ax=axes[1], edgecolor=color, facecolor=color, alpha=0.35, linewidth=1) | |
| total += len(gdf) | |
| legend_items.append(f"{feat_type}: {len(gdf)}") | |
| title = f"{total} Features" | |
| if legend_items: | |
| title += f" ({', '.join(legend_items)})" | |
| axes[1].set_title(title, fontsize=12, fontweight=600, color="#18181b", pad=12) | |
| axes[1].tick_params(labelsize=7, colors="#71717a") | |
| for ax in axes: | |
| for spine in ax.spines.values(): | |
| spine.set_color("#e4e4e7") | |
| plt.tight_layout(pad=2) | |
| out = os.path.join(os.path.dirname(geotiff), "overlay.png") | |
| plt.savefig(out, bbox_inches="tight", facecolor="#fafafa") | |
| plt.close() | |
| return out | |
| # --------------------------------------------------------------------------- | |
| # Dashboard formatting | |
| # --------------------------------------------------------------------------- | |
| class ConfidenceTracker: | |
| """Track running confidence stats per feature type.""" | |
| def __init__(self): | |
| self.scores: dict[str, list[float]] = defaultdict(list) | |
| def add(self, feat_type: str, score: float): | |
| self.scores[feat_type].append(score) | |
| def count(self, feat_type: str) -> int: | |
| return len(self.scores[feat_type]) | |
| def total(self) -> int: | |
| return sum(len(v) for v in self.scores.values()) | |
| def stats(self, feat_type: str) -> dict: | |
| s = self.scores[feat_type] | |
| if not s: | |
| return {"count": 0, "avg": 0, "min": 0, "max": 0, "low": 0} | |
| return { | |
| "count": len(s), | |
| "avg": sum(s) / len(s), | |
| "min": min(s), | |
| "max": max(s), | |
| "low": sum(1 for x in s if x < 0.5), | |
| } | |
| def format_dashboard( | |
| image_name: str, | |
| img_w: int, | |
| img_h: int, | |
| tile_idx: int, | |
| total_tiles: int, | |
| elapsed: float, | |
| tracker: ConfidenceTracker, | |
| feature_types: list[str], | |
| status: str = "Processing", | |
| ) -> str: | |
| pct = (tile_idx / total_tiles * 100) if total_tiles > 0 else 0 | |
| elapsed_str = _fmt_time(elapsed) | |
| if tile_idx > 0 and tile_idx < total_tiles: | |
| rate = elapsed / tile_idx | |
| remaining = rate * (total_tiles - tile_idx) | |
| remaining_str = _fmt_time(remaining) | |
| else: | |
| remaining_str = "--" | |
| md = f"**{status}: {image_name}** ({img_w:,} x {img_h:,} px)\n\n" | |
| md += f"Tile {tile_idx} / {total_tiles} | {pct:.0f}% | " | |
| md += f"Elapsed: {elapsed_str} | Remaining: ~{remaining_str}\n\n" | |
| md += "| Feature | Count | Avg Conf | Min Conf | Max Conf | Low (<0.5) |\n" | |
| md += "|---------|-------|----------|----------|----------|------------|\n" | |
| for ft in feature_types: | |
| st = tracker.stats(ft) | |
| if st["count"] > 0: | |
| md += ( | |
| f"| {ft} | {st['count']} | {st['avg']:.2f} | " | |
| f"{st['min']:.2f} | {st['max']:.2f} | {st['low']} |\n" | |
| ) | |
| else: | |
| md += f"| {ft} | 0 | -- | -- | -- | -- |\n" | |
| return md | |
| def format_final_summary( | |
| tracker: ConfidenceTracker, | |
| feature_types: list[str], | |
| final_counts: dict[str, int], | |
| elapsed: float, | |
| ) -> str: | |
| md = f"**Extraction Complete** | Total time: {_fmt_time(elapsed)}\n\n" | |
| md += "| Feature | Raw | Final | Avg Conf | Min | Max | Low (<0.5) |\n" | |
| md += "|---------|-----|-------|----------|-----|-----|------------|\n" | |
| for ft in feature_types: | |
| st = tracker.stats(ft) | |
| final = final_counts.get(ft, 0) | |
| if st["count"] > 0: | |
| md += ( | |
| f"| {ft} | {st['count']} | {final} | " | |
| f"{st['avg']:.2f} | {st['min']:.2f} | {st['max']:.2f} | {st['low']} |\n" | |
| ) | |
| else: | |
| md += f"| {ft} | 0 | 0 | -- | -- | -- | -- |\n" | |
| md += "\n*Adjust the confidence threshold and re-run to improve results.*" | |
| return md | |
| def _fmt_time(secs: float) -> str: | |
| if secs < 60: | |
| return f"{secs:.0f}s" | |
| m, s = divmod(int(secs), 60) | |
| return f"{m}m {s:02d}s" | |
| # --------------------------------------------------------------------------- | |
| # Main pipeline (generator for live updates) | |
| # --------------------------------------------------------------------------- | |
| PREVIEW_INTERVAL = 10 # update overlay every N tiles | |
| CHECKPOINT_INTERVAL = 100 # save downloadable files every N tiles | |
| def run_pipeline( | |
| image_file, | |
| world_file, | |
| crs: str, | |
| feature_types: list[str], | |
| confidence: float, | |
| tile_size: int, | |
| ): | |
| """Generator: yields (overlay_image, dashboard_md, files) after each tile.""" | |
| if image_file is None: | |
| raise gr.Error("Please upload an aerial image.") | |
| if not feature_types: | |
| raise gr.Error("Select at least one feature type.") | |
| start_time = time.time() | |
| tracker = ConfidenceTracker() | |
| all_features: list[dict] = [] | |
| tmp_dir = tempfile.mkdtemp(prefix="janus_") | |
| # -- Ingest -- | |
| yield None, "**Ingesting image...**", None | |
| geotiff = ingest(image_file, world_file, crs) | |
| # Model loads on first segment_tile call (inside @spaces.GPU context) | |
| yield None, "**Starting extraction...** (model loads on first tile)", None | |
| # -- Compute tiles -- | |
| with rasterio.open(geotiff) as src: | |
| img_w, img_h = src.width, src.height | |
| crs_obj = src.crs | |
| transform = src.transform | |
| image_name = Path(image_file).name | |
| overlap = tile_size // 8 | |
| windows = compute_tile_windows(img_w, img_h, tile_size, overlap) | |
| total_tiles = len(windows) | |
| dash = format_dashboard(image_name, img_w, img_h, 0, total_tiles, 0, tracker, feature_types, "Starting") | |
| yield None, dash, None | |
| # -- Process tiles -- | |
| # Track the scanned region for progress visualization | |
| scan_left = scan_bottom = float("inf") | |
| scan_right = scan_top = float("-inf") | |
| with rasterio.open(geotiff) as src: | |
| for tile_idx, window in enumerate(windows): | |
| tile_data = src.read([1, 2, 3], window=window) | |
| tile_rgb = np.transpose(tile_data, (1, 2, 0)) | |
| tile_transform = rasterio.windows.transform(window, transform) | |
| # Update scanned bounds | |
| tile_bounds = rasterio.windows.bounds(window, transform) | |
| scan_left = min(scan_left, tile_bounds[0]) | |
| scan_bottom = min(scan_bottom, tile_bounds[1]) | |
| scan_right = max(scan_right, tile_bounds[2]) | |
| scan_top = max(scan_top, tile_bounds[3]) | |
| for ft in feature_types: | |
| preset = FEATURE_PRESETS[ft] | |
| masks, scores = segment_tile(tile_rgb, preset["prompt"], confidence) | |
| if len(masks) == 0: | |
| continue | |
| labeled = np.zeros((tile_rgb.shape[0], tile_rgb.shape[1]), dtype=np.int32) | |
| for i, mask in enumerate(masks): | |
| labeled[mask > 0] = i + 1 | |
| for geom, val in shapes(labeled, mask=(labeled > 0), transform=tile_transform): | |
| val = int(val) | |
| score = float(scores[val - 1]) if val - 1 < len(scores) else 0.0 | |
| tracker.add(ft, score) | |
| all_features.append({ | |
| "geometry": shape(geom), | |
| "feature_type": ft, | |
| "confidence": round(score, 3), | |
| }) | |
| elapsed = time.time() - start_time | |
| dash = format_dashboard( | |
| image_name, img_w, img_h, | |
| tile_idx + 1, total_tiles, | |
| elapsed, tracker, feature_types, | |
| ) | |
| # Checkpoint: save downloadable files every N tiles | |
| # Update overlay + save files every PREVIEW_INTERVAL tiles | |
| if (tile_idx + 1) % PREVIEW_INTERVAL == 0 or tile_idx == total_tiles - 1 or tile_idx == 0: | |
| temp_by_type = {} | |
| for ft in feature_types: | |
| ft_feats = [f for f in all_features if f["feature_type"] == ft] | |
| if ft_feats: | |
| temp_by_type[ft] = gpd.GeoDataFrame(ft_feats, crs=crs_obj) | |
| # Save files every update so they're always downloadable | |
| checkpoint_files = None | |
| if temp_by_type: | |
| checkpoint_files = export_all(temp_by_type, tmp_dir) | |
| scanned = (scan_left, scan_bottom, scan_right, scan_top) | |
| overlay = make_overlay(geotiff, temp_by_type, scanned_bounds=scanned) | |
| yield overlay, dash, checkpoint_files | |
| else: | |
| yield gr.update(), dash, gr.update() | |
| # -- Deduplicate + filter -- | |
| elapsed = time.time() - start_time | |
| yield gr.update(), format_dashboard( | |
| image_name, img_w, img_h, total_tiles, total_tiles, | |
| elapsed, tracker, feature_types, "Deduplicating & filtering", | |
| ), None | |
| final_by_type: dict[str, gpd.GeoDataFrame] = {} | |
| final_counts: dict[str, int] = {} | |
| for ft in feature_types: | |
| ft_feats = [f for f in all_features if f["feature_type"] == ft] | |
| if not ft_feats: | |
| final_counts[ft] = 0 | |
| continue | |
| gdf = merge_and_deduplicate(ft_feats, crs_obj) | |
| gdf = filter_features(gdf, FEATURE_PRESETS[ft]) | |
| final_by_type[ft] = gdf | |
| final_counts[ft] = len(gdf) | |
| # -- Export -- | |
| export_paths = export_all(final_by_type, tmp_dir) | |
| # -- Final overlay -- | |
| final_overlay = make_overlay(geotiff, final_by_type) | |
| elapsed = time.time() - start_time | |
| final_dash = format_final_summary(tracker, feature_types, final_counts, elapsed) | |
| yield final_overlay, final_dash, export_paths | |
| # --------------------------------------------------------------------------- | |
| # Gradio UI | |
| # --------------------------------------------------------------------------- | |
| CUSTOM_CSS = """ | |
| @import url('https://fonts.googleapis.com/css2?family=Outfit:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap'); | |
| :root { | |
| --body-background-fill: #fafafa !important; | |
| --block-background-fill: #ffffff !important; | |
| --block-border-color: #e4e4e7 !important; | |
| --block-label-text-color: #3f3f46 !important; | |
| --block-title-text-color: #18181b !important; | |
| --button-primary-background-fill: #18181b !important; | |
| --button-primary-text-color: #fafafa !important; | |
| --button-primary-background-fill-hover: #27272a !important; | |
| --input-background-fill: #ffffff !important; | |
| --border-color-primary: #e4e4e7 !important; | |
| } | |
| * { font-family: 'Outfit', system-ui, -apple-system, sans-serif !important; } | |
| code, pre, .code, [class*="mono"] { font-family: 'JetBrains Mono', monospace !important; } | |
| .gradio-container { max-width: 1400px !important; margin: 0 auto !important; background: #fafafa !important; } | |
| .gr-button-primary { | |
| border-radius: 12px !important; font-weight: 600 !important; | |
| letter-spacing: -0.01em !important; padding: 12px 32px !important; | |
| transition: all 0.2s cubic-bezier(0.16, 1, 0.3, 1) !important; | |
| box-shadow: 0 1px 3px rgba(0,0,0,0.08) !important; | |
| } | |
| .gr-button-primary:hover { transform: translateY(-1px) !important; box-shadow: 0 4px 12px rgba(0,0,0,0.1) !important; } | |
| .gr-button-primary:active { transform: translateY(0) scale(0.98) !important; } | |
| .gr-panel, .gr-box, .gr-form { | |
| border-radius: 20px !important; border: 1px solid #e4e4e7 !important; | |
| box-shadow: 0 20px 40px -15px rgba(0,0,0,0.04) !important; | |
| } | |
| .gr-input, .gr-textbox textarea, select { | |
| border-radius: 12px !important; border: 1px solid #e4e4e7 !important; | |
| font-size: 14px !important; transition: border-color 0.2s ease !important; | |
| } | |
| .gr-input:focus, .gr-textbox textarea:focus { | |
| border-color: #18181b !important; box-shadow: 0 0 0 3px rgba(24,24,27,0.06) !important; | |
| } | |
| h1 { font-size: 2.25rem !important; font-weight: 700 !important; letter-spacing: -0.03em !important; line-height: 1.1 !important; color: #18181b !important; } | |
| h2, h3, .gr-block-label { font-weight: 600 !important; letter-spacing: -0.02em !important; color: #3f3f46 !important; } | |
| .markdown-text p { color: #52525b !important; line-height: 1.6 !important; max-width: 65ch !important; } | |
| footer { display: none !important; } | |
| """ | |
| HEADER_MD = """ | |
| # Janus | |
| ### GIS-ready feature extraction from aerial imagery | |
| Upload a GeoTIFF (any size), select feature types to extract, and watch SAM 3 process | |
| tile by tile with live confidence tracking. | |
| """ | |
| def _gpu_status(): | |
| try: | |
| if torch.cuda.is_available(): | |
| return ( | |
| f"Running on **{torch.cuda.get_device_name(0)}** " | |
| f"({torch.cuda.get_device_properties(0).total_memory / (1024**3):.0f} GB)" | |
| ) | |
| except Exception: | |
| pass | |
| return "No GPU detected — inference will be slow" | |
| gpu_status = _gpu_status() | |
| with gr.Blocks(css=CUSTOM_CSS, title="Janus — Feature Extraction") as demo: | |
| gr.Markdown(HEADER_MD) | |
| gr.Markdown(f"*{gpu_status}*") | |
| with gr.Row(equal_height=False): | |
| with gr.Column(scale=3): | |
| with gr.Group(): | |
| image_input = gr.File( | |
| label="Aerial Image", | |
| file_types=[".png", ".jpg", ".jpeg", ".tif", ".tiff"], | |
| type="filepath", | |
| ) | |
| world_input = gr.File( | |
| label="World File (optional for GeoTIFF)", | |
| file_types=[".pgw", ".jgw", ".tfw"], | |
| type="filepath", | |
| ) | |
| crs_input = gr.Textbox(value="EPSG:4326", label="CRS") | |
| feature_checks = gr.CheckboxGroup( | |
| choices=list(FEATURE_PRESETS.keys()), | |
| value=["Building"], | |
| label="Feature Types", | |
| info="Select one or more feature types to extract", | |
| ) | |
| with gr.Row(): | |
| confidence_slider = gr.Slider( | |
| minimum=0.1, maximum=0.95, value=0.5, step=0.05, | |
| label="Confidence Threshold", | |
| ) | |
| tile_size_dropdown = gr.Dropdown( | |
| choices=[256, 512, 1024, 2048], | |
| value=256, | |
| label="Tile Size (px)", | |
| info="Larger = fewer tiles but more GPU memory", | |
| ) | |
| run_btn = gr.Button("Extract Features", variant="primary", size="lg") | |
| with gr.Column(scale=5): | |
| dashboard = gr.Markdown(value="*Upload an image and click Extract Features to begin.*") | |
| output_image = gr.Image(label="Live Preview", type="filepath", show_download_button=True) | |
| file_output = gr.Files(label="Download GIS Files") | |
| run_btn.click( | |
| fn=run_pipeline, | |
| inputs=[image_input, world_input, crs_input, feature_checks, confidence_slider, tile_size_dropdown], | |
| outputs=[output_image, dashboard, file_output], | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |