import spaces # noqa: F401 (must precede any CUDA-touching import) import numpy as np from PIL import Image import gradio as gr import torch import torchvision.transforms as transforms from huggingface_hub import hf_hub_download from ultralytics import YOLO from cabinet import CABiNet IMGSZ = 1024 # --- UAVid Model Zoo --------------------------------------------------------- # Metrics/params/FLOPs sourced from hf_modelcards/model_metrics/*/metrics.json # and hf_modelcards/generate_hf_model_zoo.py's MODEL_INFO (kept in sync with # the main README's UAVid Model Zoo table). Ordered by mIoU descending, so the # first entry (best model) is the dropdown's default. MODEL_ZOO = { "cabinet-mobilenetv3-large": dict( display_name="CABiNet (MobileNetV3-Large)", framework="cabinet", cabinet_mode="large", repo_id="dronefreak/cabinet-mobilenetv3-large-uavid", weights_filename="cabinet_best.pth", miou=68.60, params_m=9.17, flops_g=54.8, ), "cabinet-mobilenetv3-small": dict( display_name="CABiNet (MobileNetV3-Small)", framework="cabinet", cabinet_mode="small", repo_id="dronefreak/cabinet-mobilenetv3-small-uavid", weights_filename="cabinet_best.pth", miou=66.84, params_m=5.36, flops_g=44.1, ), "yolo26x-sem": dict( display_name="YOLO26x-sem", framework="ultralytics", cabinet_mode=None, repo_id="dronefreak/uavid-yolo26x-sem", weights_filename="best.pt", miou=64.41, params_m=40.16, flops_g=430.9, ), "yolo26l-sem": dict( display_name="YOLO26l-sem", framework="ultralytics", cabinet_mode=None, repo_id="dronefreak/uavid-yolo26l-sem", weights_filename="best.pt", miou=63.28, params_m=17.87, flops_g=192.4, ), "yolo26m-sem": dict( display_name="YOLO26m-sem", framework="ultralytics", cabinet_mode=None, repo_id="dronefreak/uavid-yolo26m-sem", weights_filename="best.pt", miou=61.98, params_m=14.32, flops_g=152.3, ), "yolo26s-sem": dict( display_name="YOLO26s-sem", framework="ultralytics", cabinet_mode=None, repo_id="dronefreak/uavid-yolo26s-sem", weights_filename="best.pt", miou=61.69, params_m=6.50, flops_g=44.4, ), "yolo26n-sem": dict( display_name="YOLO26n-sem", framework="ultralytics", cabinet_mode=None, repo_id="dronefreak/uavid-yolo26n-sem", weights_filename="best.pt", miou=58.17, params_m=1.63, flops_g=11.4, ), } DEFAULT_MODEL_SLUG = next(iter(MODEL_ZOO)) # best mIoU: cabinet-mobilenetv3-large # k, t, c, SE, HS, s — mirrors configs/model/mobilenetv3_{large,small}.yaml in the # CABiNet training repo (https://github.com/dronefreak/CABiNet). CABINET_CFGS = { "large": [ [3, 1, 16, 0, 0, 1], [3, 4, 24, 0, 0, 2], [3, 3, 24, 0, 0, 1], [5, 3, 40, 1, 0, 2], [5, 3, 40, 1, 0, 1], [5, 3, 40, 1, 0, 1], [3, 6, 80, 0, 1, 2], [3, 2.5, 80, 0, 1, 1], [3, 2.3, 80, 0, 1, 1], [3, 2.3, 80, 0, 1, 1], [3, 6, 112, 1, 1, 1], [3, 6, 112, 1, 1, 1], [5, 6, 160, 1, 1, 2], [5, 6, 160, 1, 1, 1], [5, 6, 160, 1, 1, 1], ], "small": [ [3, 1, 16, 1, 0, 2], [3, 4.5, 24, 0, 0, 2], [3, 3.67, 24, 0, 0, 1], [5, 4, 40, 1, 1, 2], [5, 6, 40, 1, 1, 1], [5, 6, 40, 1, 1, 1], [5, 3, 48, 1, 1, 1], [5, 3, 48, 1, 1, 1], [5, 6, 96, 1, 1, 2], [5, 6, 96, 1, 1, 1], [5, 6, 96, 1, 1, 1], ], } # Mean/std computed from the UAVid train set (see # src/datasets/compute_uavid_stats.py in the CABiNet training repo). CABINET_TO_TENSOR = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize(mean=(0.480, 0.499, 0.457), std=(0.225, 0.208, 0.228)), ] ) # --- UAVid classes + official-style color palette (RGB) --- CLASS_NAMES = { 0: "Clutter", 1: "Building", 2: "Road", 3: "Static Car", 4: "Tree", 5: "Vegetation", 6: "Human", 7: "Moving Car", } PALETTE = np.array( [ [0, 0, 0], # 0 Clutter - black [128, 0, 0], # 1 Building - dark red [128, 64, 128], # 2 Road - purple/grey [192, 0, 192], # 3 Static Car - magenta [0, 128, 0], # 4 Tree - green [128, 128, 0], # 5 Vegetation - olive [64, 64, 0], # 6 Human - dark yellow [64, 0, 128], # 7 Moving Car - blue/purple ], dtype=np.uint8, ) # --- Model cache ------------------------------------------------------------- # Models are downloaded/instantiated lazily (on first selection) rather than # all at module scope, since eagerly loading all 7 checkpoints would slow # cold start; the cache dict itself lives at module scope so a model is only # ever downloaded/built once per running Space instance. _MODEL_CACHE: dict = {} def _get_model(slug: str): if slug in _MODEL_CACHE: return _MODEL_CACHE[slug] info = MODEL_ZOO[slug] weights = hf_hub_download(repo_id=info["repo_id"], filename=info["weights_filename"]) if info["framework"] == "ultralytics": model = YOLO(weights) else: cfgs = CABINET_CFGS[info["cabinet_mode"]] model = CABiNet(n_classes=len(CLASS_NAMES), cfgs=cfgs, mode=info["cabinet_mode"]) ckpt = torch.load(weights, map_location="cpu", weights_only=True) state_dict = ( ckpt["model_state"] if isinstance(ckpt, dict) and "model_state" in ckpt else ckpt ) model.load_state_dict(state_dict) model.eval() try: model.to("cuda") except Exception as e: # pragma: no cover - only relevant off-GPU print(f"model.to('cuda') deferred for {slug}: {e!r}") _MODEL_CACHE[slug] = model return model def _predict_mask(slug: str, pil: Image.Image) -> np.ndarray: """Run inference and return a (H, W) class-ID mask at the model's native output size.""" info = MODEL_ZOO[slug] model = _get_model(slug) if info["framework"] == "ultralytics": results = model.predict(source=pil, task="semantic", imgsz=IMGSZ, verbose=False) mask = results[0].semantic_mask.data if hasattr(mask, "cpu"): mask = mask.cpu().numpy() return np.asarray(mask).astype(np.int64) # CABiNet is fully convolutional but expects a fixed-size batch of one; # resize to a square, run a single forward pass, and upsample the # resulting class-ID mask (nearest-neighbor, not the RGB image) back up. resized = pil.resize((IMGSZ, IMGSZ), Image.BILINEAR) tensor = CABINET_TO_TENSOR(resized).unsqueeze(0) tensor = tensor.to(next(model.parameters()).device) with torch.no_grad(): logits = model(tensor)[0] return logits.argmax(dim=1).squeeze(0).cpu().numpy().astype(np.int64) def _resize_mask_to(mask: np.ndarray, size_wh: tuple) -> np.ndarray: """Resize a (H, W) class-ID mask to (W, H) = size_wh using nearest-neighbor.""" if (mask.shape[1], mask.shape[0]) == size_wh: return mask resized = Image.fromarray(mask.astype(np.uint8)).resize(size_wh, Image.NEAREST) return np.asarray(resized).astype(np.int64) def _colorize(mask: np.ndarray) -> np.ndarray: """Map a (H, W) class-ID mask to an (H, W, 3) RGB image.""" return PALETTE[np.clip(mask, 0, len(PALETTE) - 1)] def _legend_html(present_ids, class_pct) -> str: rows = [] for cid in present_ids: r, g, b = PALETTE[cid] pct = class_pct.get(cid, 0.0) rows.append( f'
' f'' f'{CLASS_NAMES[cid]} — {pct:.1f}%
' ) return ( '
' "Detected classes" + "".join(rows) + "
" ) @spaces.GPU(duration=60) def segment( image: Image.Image, model_choice: str = DEFAULT_MODEL_SLUG, overlay_opacity: float = 0.55, ): """Run UAVid semantic segmentation on an aerial/drone image. Args: image: Input RGB image (aerial / oblique drone view of an urban scene). model_choice: Key into MODEL_ZOO selecting which model to run. overlay_opacity: Blend factor for the colored mask over the photo (0-1). Returns: Tuple of (overlay image, class-color legend HTML). """ if image is None: raise gr.Error("Please provide an input image.") if model_choice not in MODEL_ZOO: raise gr.Error(f"Unknown model: {model_choice}") pil = image.convert("RGB") mask = _predict_mask(model_choice, pil) mask = _resize_mask_to(mask, pil.size) color_mask = _colorize(mask) base = np.asarray(pil).astype(np.float32) a = float(np.clip(overlay_opacity, 0.0, 1.0)) blended = (base * (1.0 - a) + color_mask.astype(np.float32) * a).astype(np.uint8) ids, counts = np.unique(mask, return_counts=True) total = mask.size class_pct = {int(i): 100.0 * c / total for i, c in zip(ids, counts)} present = [int(i) for i in ids if 0 <= int(i) < len(PALETTE)] return Image.fromarray(blended), _legend_html(present, class_pct) CSS = """ #col-container { max-width: 1100px; margin: 0 auto; } .dark .gradio-container { color: var(--body-text-color); } """ with gr.Blocks() as demo: with gr.Column(elem_id="col-container"): gr.Markdown( """ # UAVid Semantic Segmentation — CABiNet & YOLO26 Model Zoo Pick any model from the [**UAVid Semantic Segmentation Model Zoo**](https://huggingface.co/collections/dronefreak/uavid-semantic-segmentation-model-zoo) — CABiNet (custom dual-branch PyTorch) or Ultralytics YOLO26 (n/s/m/l/x) — all fine-tuned on [UAVid](https://uavid.nl/) to label every pixel of an oblique drone / aerial urban scene into 8 classes: Clutter, Building, Road, Static Car, Tree, Vegetation, Human, Moving Car. **CABiNet (MobileNetV3-Large) is the top performer** — it beats every YOLO26 variant, including the largest (YOLO26x), on mIoU while using a fraction of the compute. Upload an aerial street-scene image (or try an example) to get a colored segmentation overlay. """ ) with gr.Row(): with gr.Column(): inp = gr.Image(type="pil", label="Input aerial image") model_dd = gr.Dropdown( choices=[ ( f"{v['display_name']} — {v['miou']:.2f}% mIoU · {v['flops_g']:.1f} GFLOPs", k, ) for k, v in MODEL_ZOO.items() ], value=DEFAULT_MODEL_SLUG, label="Model", ) run = gr.Button("Segment", variant="primary") with gr.Accordion("Advanced settings", open=False): opacity = gr.Slider( 0.0, 1.0, value=0.55, step=0.05, label="Overlay opacity", ) with gr.Column(): out_img = gr.Image(type="pil", label="Segmentation overlay") out_legend = gr.HTML(label="Legend") run.click( fn=segment, inputs=[inp, model_dd, opacity], outputs=[out_img, out_legend], api_name="segment", ) gr.Examples( examples=[ ["examples/uavid_1.jpg"], ["examples/uavid_2.jpg"], ["examples/uavid_3.jpg"], ], inputs=[inp], outputs=[out_img, out_legend], fn=segment, cache_examples=True, cache_mode="lazy", ) if __name__ == "__main__": demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)