Spaces:
Running on Zero
Running on Zero
| 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'<div style="display:flex;align-items:center;gap:8px;margin:2px 0;">' | |
| f'<span style="display:inline-block;width:16px;height:16px;border-radius:3px;' | |
| f'background:rgb({r},{g},{b});border:1px solid rgba(128,128,128,.5);"></span>' | |
| f'<span>{CLASS_NAMES[cid]} — {pct:.1f}%</span></div>' | |
| ) | |
| return ( | |
| '<div style="font-size:14px;">' | |
| "<b>Detected classes</b>" + "".join(rows) + "</div>" | |
| ) | |
| 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) | |