Spaces:
Running on Zero
Running on Zero
| import os | |
| import sys | |
| from pathlib import Path | |
| APP_DIR = Path(__file__).resolve().parent | |
| DATA_DIR = APP_DIR / "urbex_data" | |
| CITY_ROOT = DATA_DIR / "cityrefer" | |
| CACHE_ROOT = CITY_ROOT / "verify_region_cache" / "cityrefer" | |
| # The released environment reads these at import time (see notebooks/urbex_colab.ipynb). | |
| os.environ["CITY_AGENT_DATA"] = str(CITY_ROOT) | |
| os.environ["CITY_AGENT_DATA_STPLS3D"] = str(DATA_DIR / "stpls3d") | |
| os.environ["CITY_TOOL_REQUIRE_CACHE"] = "1" | |
| sys.path.insert(0, str(APP_DIR)) | |
| import spaces # noqa: E402 (must precede torch) | |
| import torch # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import json # noqa: E402 | |
| import tempfile # noqa: E402 | |
| import time # noqa: E402 | |
| from concurrent.futures import ThreadPoolExecutor # noqa: E402 | |
| from PIL import Image, ImageDraw # noqa: E402 | |
| from transformers import AutoProcessor, Qwen3VLForConditionalGeneration # noqa: E402 | |
| from scripts.data.hub_assets import fetch, load_manifest, unpack # noqa: E402 | |
| MODEL_ID = "vvangfaye/Urbex-4B" | |
| DATASET_ID = "vvangfaye/Urbex-Data" | |
| MAX_NEW_TOKENS = 1024 | |
| CONTEXT_LIMIT = 16384 | |
| # -------------------------------------------------------------------------- | |
| # 1. Download the released CityRefer test scenes (BEV + cached oblique views) | |
| # -------------------------------------------------------------------------- | |
| t0 = time.time() | |
| manifest, revision = load_manifest(DATASET_ID, DATA_DIR) | |
| city = manifest["datasets"]["cityrefer"] | |
| for entry in city["files"]: | |
| fetch(DATASET_ID, entry, DATA_DIR, revision) | |
| SCENES = sorted(city["scenes"]) | |
| def _download_scene(scene): | |
| entry = city["scenes"][scene] | |
| for file in entry["files"]: | |
| fetch(DATASET_ID, file, DATA_DIR, revision) | |
| if (DATA_DIR / entry["extract_prefix"] / ".urbex-download.json").is_file(): | |
| return scene | |
| archive = fetch(DATASET_ID, entry["archive"], DATA_DIR, revision, install=False) | |
| unpack(archive, DATA_DIR, entry["extract_prefix"], entry["archive"]) | |
| try: # the extracted copy is all we need | |
| Path(archive).resolve().unlink() | |
| except OSError: | |
| pass | |
| return scene | |
| with ThreadPoolExecutor(max_workers=6) as pool: | |
| for s in pool.map(_download_scene, SCENES): | |
| print("scene ready:", s, flush=True) | |
| print(f"Data ready in {time.time() - t0:.0f}s", flush=True) | |
| from roll.pipeline.agentic.env.city_grounding import CityGroundingEnv # noqa: E402 | |
| from roll.pipeline.agentic.env.city_grounding.render import _load_bev_cached # noqa: E402 | |
| from scripts.demo.view_3d import rebuild_scene, viewer_html # noqa: E402 | |
| # Released test annotations: used for scene-wise example queries and to report | |
| # BEV IoU against ground truth when the query matches a benchmark sample. | |
| GT = {} | |
| for split in ("test_ND", "test_NO"): | |
| for line in (CITY_ROOT / "meta" / f"cityrefer_{split}.jsonl").read_text().splitlines(): | |
| if line.strip(): | |
| s = json.loads(line) | |
| GT[(s["scene_id"], s["query"].strip())] = s["gt_bbox"] | |
| def scene_label(scene): | |
| town, _, idx = scene.partition("_block_") | |
| return f"{town.capitalize()} · block {idx}" | |
| SCENE_CHOICES = [(scene_label(s), s) for s in SCENES] | |
| # Small previews of each global BEV (the full PNGs are ~100 MB each). | |
| PREVIEW_DIR = DATA_DIR / "previews" | |
| PREVIEW_DIR.mkdir(exist_ok=True) | |
| MAPPINGS = {} | |
| for s in SCENES: | |
| png = CITY_ROOT / "bev_cityrefer" / f"{s}.png" | |
| MAPPINGS[s] = json.loads(png.with_suffix(".json").read_text()) | |
| out = PREVIEW_DIR / f"{s}.jpg" | |
| if not out.exists(): | |
| with Image.open(png) as im: | |
| im.draft("RGB", (2048, 2048)) | |
| im = im.convert("RGB") | |
| im.thumbnail((2048, 2048), Image.Resampling.LANCZOS) | |
| im.save(out, quality=88) | |
| print("previews ready", flush=True) | |
| def preview(scene): | |
| return Image.open(PREVIEW_DIR / f"{scene}.jpg").convert("RGB") | |
| # -------------------------------------------------------------------------- | |
| # 2. Load the model (bf16, full precision, eagerly on GPU) | |
| # -------------------------------------------------------------------------- | |
| processor = AutoProcessor.from_pretrained(MODEL_ID) | |
| model = Qwen3VLForConditionalGeneration.from_pretrained( | |
| MODEL_ID, dtype=torch.bfloat16, attn_implementation="sdpa" | |
| ).to("cuda").eval() | |
| def collect_images(conversation): | |
| return [item["image"] for message in conversation | |
| if isinstance(message.get("content"), list) | |
| for item in message["content"] | |
| if isinstance(item, dict) and item.get("type") == "image"] | |
| def _fence(text): | |
| return "```text\n" + text.replace("```", "'''").strip() + "\n```" | |
| def _draw_norm_box(img, bbox, color=(255, 30, 30), width=4): | |
| img = img.copy() | |
| w, h = img.size | |
| x0, y0, x1, y1 = [float(v) / 1000.0 for v in bbox] | |
| ImageDraw.Draw(img).rectangle([x0 * w, y0 * h, x1 * w, y1 * h], outline=color, width=width) | |
| return img | |
| def _draw_world_boxes(img, mapping, boxes): | |
| """Draw world-XY boxes [(box, colour, width), ...] onto a (downscaled) global BEV.""" | |
| img = img.copy() | |
| W, H = img.size | |
| sx = W / (mapping["x_max"] - mapping["x_min"]) | |
| sy = H / (mapping["y_max"] - mapping["y_min"]) | |
| d = ImageDraw.Draw(img) | |
| for (x0, y0, x1, y1), colour, width in boxes: | |
| d.rectangle([(x0 - mapping["x_min"]) * sx, (mapping["y_max"] - y1) * sy, | |
| (x1 - mapping["x_min"]) * sx, (mapping["y_max"] - y0) * sy], | |
| outline=colour, width=width) | |
| return img | |
| # -------------------------------------------------------------------------- | |
| # 3. Agent loop (matches the reference Colab: greedy, 8 turns, 1024 tokens) | |
| # -------------------------------------------------------------------------- | |
| def _gpu_duration(scene_id, query, max_turns, *args, **kwargs): | |
| # Measured on the live Space: ~4-5 s per agent turn (4-turn runs took 14-19 s). | |
| return 10 + 6 * int(max_turns) | |
| def _agent_gpu(scene_id, query, max_turns): | |
| sample = {"dataset": "cityrefer", "scene_id": scene_id, "query": query, | |
| "bev_png": f"bev_cityrefer/{scene_id}.png"} | |
| gt = GT.get((scene_id, query.strip())) | |
| if gt is not None: | |
| sample["gt_bbox"] = gt | |
| with tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False) as f: | |
| f.write(json.dumps(sample) + "\n") | |
| data_path = f.name | |
| env = CityGroundingEnv( | |
| data_path=data_path, mode="val", seed=42, multiview_cache_dir=str(CACHE_ROOT), | |
| max_tool_calls=max_turns - 1, max_assistant_turns=max_turns, | |
| display_max_side=768, bbox_norm_max=1000.0, | |
| ) | |
| messages, reset_info = env.reset(seed=0) | |
| os.unlink(data_path) | |
| tools = reset_info["tools"] | |
| torch.manual_seed(42) | |
| log = [f"**Query:** {query}\n\nThe agent starts from the global bird's-eye view of `{scene_id}`."] | |
| gallery = [(env._env.view_image("global"), "global BEV (initial observation)")] | |
| final_info = {} | |
| t_start = time.perf_counter() | |
| yield {"log": "\n\n".join(log), "gallery": list(gallery), "done": False} | |
| for turn in range(1, max_turns + 1): | |
| prompt = processor.apply_chat_template( | |
| messages, tools=tools, tokenize=False, add_generation_prompt=True) | |
| images = collect_images(messages) | |
| inputs = processor(text=[prompt], images=images or None, padding=True, | |
| return_tensors="pt").to(model.device) | |
| n_prompt = inputs.input_ids.shape[1] | |
| if n_prompt + MAX_NEW_TOKENS > CONTEXT_LIMIT: | |
| log.append("⚠️ Conversation exceeded the released 16K context; stopping.") | |
| break | |
| with torch.inference_mode(): | |
| out = model.generate(**inputs, max_new_tokens=MAX_NEW_TOKENS, | |
| do_sample=False, use_cache=True) | |
| action = processor.tokenizer.decode(out[0, n_prompt:], skip_special_tokens=False) | |
| del inputs, out | |
| shown = action.replace("<|im_end|>", "").replace("<|endoftext|>", "") | |
| log.append(f"### Turn {turn} · model\n" + _fence(shown)) | |
| prior = len(messages) | |
| messages, _reward, terminated, truncated, final_info = env.step(action) | |
| for m in messages[prior:]: | |
| if m.get("role") != "tool": | |
| continue | |
| content = m.get("content", []) | |
| name = m.get("name", "tool") | |
| if isinstance(content, str): | |
| log.append(f"**Tool ({name}):**\n" + _fence(content)) | |
| continue | |
| for item in content: | |
| if item.get("type") == "text": | |
| log.append(f"**Tool ({name}):**\n" + _fence(item["text"])) | |
| elif item.get("type") == "image": | |
| caption = f"turn {turn} · {name}" | |
| if final_info.get("view_id") and name == "zoom_in": | |
| caption += f" → view {final_info['view_id']}" | |
| elif final_info.get("direction"): | |
| caption += f" ({final_info['direction']})" | |
| gallery.append((item["image"], caption)) | |
| yield {"log": "\n\n".join(log), "gallery": list(gallery), "done": False} | |
| if terminated or truncated: | |
| break | |
| elapsed = time.perf_counter() - t_start | |
| answer = final_info.get("city_answer") | |
| world_bbox = final_info.get("city_world_bbox") | |
| result = {"scene_id": scene_id, "query": query, "answer": answer, | |
| "world_bbox_xyxy_m": world_bbox, "turns": env._assistant_turns, | |
| "agent_seconds": round(elapsed, 1)} | |
| zooms = [] | |
| if answer and world_bbox: | |
| view_id = answer.get("view_id", "global") | |
| view = env._env.view_registry.get(view_id) | |
| if view is not None: | |
| gallery.append((_draw_norm_box(view["image"], answer["bbox_2d"]), | |
| f"final answer on view {view_id}")) | |
| vid = view_id | |
| while vid and vid != "global" and vid in env._env.view_registry: | |
| v = env._env.view_registry[vid] | |
| zooms.insert(0, v["world_bounds"]) | |
| vid = v["parent_id"] | |
| if gt is not None: | |
| result["bev_iou_vs_ground_truth"] = round(float(final_info.get("city_iou_2d", 0.0)), 4) | |
| else: | |
| result["reason"] = final_info.get("city_reason", "turn limit reached") | |
| print(f"[urbex] scene={scene_id} turns={result['turns']} agent_seconds={elapsed:.1f}", flush=True) | |
| yield {"log": "\n\n".join(log), "gallery": list(gallery), "done": True, | |
| "result": result, "zooms": zooms} | |
| def ground(scene_id: str, query: str, max_turns: int = 8): | |
| """Ground a natural-language description to an object in a city-scale 3D scene. | |
| The Urbex-4B agent looks at the scene's bird's-eye view, then actively searches it by | |
| looking up landmarks, zooming into regions and inspecting cached oblique 3D views, and | |
| finally commits a box. | |
| Args: | |
| scene_id: CityRefer/SensatUrban scene id, e.g. "cambridge_block_3". | |
| query: description of the object (car, building, parking lot, ground area...). | |
| max_turns: maximum assistant turns (the last turn is reserved for the answer). | |
| Returns: | |
| Search trajectory (markdown), gallery of observed views, global BEV with the | |
| prediction, interactive 3D view (HTML) and a JSON result with the world-space box. | |
| """ | |
| if scene_id not in MAPPINGS: | |
| raise gr.Error(f"Unknown scene {scene_id!r}") | |
| query = (query or "").strip() | |
| if not query: | |
| raise gr.Error("Please describe the object to find.") | |
| max_turns = int(max(2, min(8, max_turns))) | |
| bev = preview(scene_id) | |
| placeholder = "<div style='padding:1em;opacity:.7'>The 3D view appears once the agent commits an answer.</div>" | |
| state = None | |
| for state in _agent_gpu(scene_id, query, max_turns): | |
| if not state["done"]: | |
| yield state["log"], state["gallery"], bev, placeholder, {} | |
| result = state["result"] | |
| wb = result.get("world_bbox_xyxy_m") | |
| overview, html3d = bev, "<div style='padding:1em'>No valid final prediction — try rephrasing the query or using more turns.</div>" | |
| if wb: | |
| boxes = [(z, (255, 200, 0), 3) for z in state["zooms"]] + [(wb, (255, 30, 30), 6)] | |
| overview = _draw_world_boxes(bev, MAPPINGS[scene_id], boxes) | |
| try: | |
| scene = rebuild_scene(CACHE_ROOT, scene_id, wb, max_points=200_000) | |
| html3d = viewer_html(scene, title="Urbex prediction", subtitle=scene_label(scene_id), | |
| query=query, bev=(bev, MAPPINGS[scene_id]), zooms=state["zooms"], | |
| height=520) | |
| except Exception as e: # e.g. a box outside the cached area | |
| html3d = f"<div style='padding:1em'>3D view unavailable: {e}</div>" | |
| yield state["log"], state["gallery"], overview, html3d, result | |
| # -------------------------------------------------------------------------- | |
| # 4. UI | |
| # -------------------------------------------------------------------------- | |
| EXAMPLES = [ | |
| ["cambridge_block_3", "A red van that has a black car behind it, on Castle Road that is a cross from a bus stop marking on the other side of that road"], | |
| ["birmingham_block_12", "A square building with a black flat roof surrounded by cars. It has a big car parked on its left, right and front. It is to the right of Holford drive road."], | |
| ["birmingham_block_9", "The large gray triangular parking lot next to the large green yard on Aldridge Road."], | |
| ] | |
| CSS = """ | |
| #col-container { max-width: 1280px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="Urbex: Agentic Spatial Grounding") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# 🏙️ Urbex: Agentic Spatial Grounding for City-Scale 3D Scenes\n" | |
| "Describe an object in a city-scale 3D scene. The **Urbex-4B** agent searches the scene " | |
| "actively — landmark lookup, BEV zoom-ins and oblique 3D inspections — before committing a box. \n" | |
| "[Paper](https://huggingface.co/papers/2610.18356) · [Code](https://github.com/AMAP-ML/Urbex) · " | |
| "[Model](https://huggingface.co/vvangfaye/Urbex-4B) · [Data](https://huggingface.co/datasets/vvangfaye/Urbex-Data)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| scene = gr.Dropdown(SCENE_CHOICES, value="cambridge_block_3", label="Scene (CityRefer / SensatUrban)") | |
| query = gr.Textbox(label="What should the agent find?", lines=3, | |
| placeholder="e.g. The white car parked next to the red-roofed house on Castle Road") | |
| run = gr.Button("Find it", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| max_turns = gr.Slider(2, 8, value=8, step=1, label="Max assistant turns", | |
| info="The last turn is reserved for the final answer (paper default: 8).") | |
| with gr.Column(scale=3): | |
| overview = gr.Image(value=preview("cambridge_block_3"), label="Global BEV (red = prediction, yellow = zooms)", | |
| type="pil", height=420) | |
| view3d = gr.HTML(label="3D view of the prediction") | |
| with gr.Row(): | |
| with gr.Column(scale=3): | |
| gallery = gr.Gallery(label="What the agent saw", columns=4, height=360, object_fit="contain") | |
| with gr.Column(scale=2): | |
| result = gr.JSON(label="Result") | |
| with gr.Accordion("Search trajectory", open=True): | |
| log = gr.Markdown() | |
| gr.Examples(examples=EXAMPLES, inputs=[scene, query], | |
| outputs=[log, gallery, overview, view3d, result], | |
| fn=ground, cache_examples=True, cache_mode="lazy") | |
| scene.change(preview, inputs=scene, outputs=overview, api_visibility="private") | |
| run.click(ground, inputs=[scene, query, max_turns], | |
| outputs=[log, gallery, overview, view3d, result], api_name="ground") | |
| query.submit(ground, inputs=[scene, query, max_turns], | |
| outputs=[log, gallery, overview, view3d, result], api_visibility="private") | |
| if __name__ == "__main__": | |
| demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |