Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """3D HAMSTER — 3D Trajectory Prediction demo (ZeroGPU). | |
| Depth-aware VLM planner: predicts metric 3D end-effector trajectories (and 2D | |
| trajectories / pointing / bbox / VQA) from a single RGB image + metric depth map | |
| + a language instruction. | |
| Adapted from the reference Gradio script in the official repo | |
| (scripts/trajectory_prediction_gradio.py) to run on ZeroGPU: the model is loaded | |
| once at module scope onto CUDA and inference is wrapped in @spaces.GPU. | |
| """ | |
| import os | |
| # DINOv2 geometry encoder uses xformers when available; force the pure-torch | |
| # fallback so we don't need an xformers CUDA build on Blackwell. | |
| os.environ.setdefault("XFORMERS_DISABLED", "1") | |
| import spaces # noqa: E402 MUST come before torch / transformers | |
| import json # noqa: E402 | |
| import re # noqa: E402 | |
| import tempfile # noqa: E402 | |
| from pathlib import Path # noqa: E402 | |
| from typing import Optional # noqa: E402 | |
| import cv2 # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import plotly.graph_objects as go # noqa: E402 | |
| import torch # noqa: E402 | |
| from PIL import Image # noqa: E402 | |
| from huggingface_hub import snapshot_download # noqa: E402 | |
| # ── Constants (mirror the reference script) ────────────────────────────────── | |
| MODEL_ID = "DAVIAN-Robotics/3D_HAMSTER" | |
| HERE = Path(__file__).parent.resolve() | |
| EXAMPLES_DIR = str(HERE / "examples") | |
| TARGET_SIZE = 640 # training resolution: longest edge = 640 | |
| V5_SYSTEM_PROMPT = "" | |
| VQA_STYLE = "General VQA" | |
| BBOX_STYLE = "2D Bounding Box" | |
| V5_PROMPT_SUFFIXES = { | |
| "3D Trajectory": ( | |
| "Predict the full manipulation trajectory as point_3d waypoints " | |
| "with depth and gripper state in JSON." | |
| ), | |
| "2D Trajectory": ( | |
| "Predict the full manipulation trajectory as point_2d waypoints " | |
| "with gripper state in JSON." | |
| ), | |
| "3D Pointing": "Report the point_3d location in JSON.", | |
| "2D Pointing": "Report point_2d locations in JSON.", | |
| BBOX_STYLE: None, | |
| VQA_STYLE: None, | |
| } | |
| V5_PROMPT_STYLES = list(V5_PROMPT_SUFFIXES.keys()) | |
| V5_DEFAULT_STYLE = "3D Trajectory" | |
| POINTING_STYLES = {"2D Pointing", "3D Pointing"} | |
| def build_v5_human_message(instruction: str, prompt_style: str) -> str: | |
| instr = instruction.strip() | |
| if prompt_style == BBOX_STYLE: | |
| return ( | |
| f"I'm looking for {instr} in this image. Can you locate it? " | |
| "Report bbox coordinates in JSON format." | |
| ) | |
| if prompt_style not in V5_PROMPT_SUFFIXES: | |
| prompt_style = V5_DEFAULT_STYLE | |
| suffix = V5_PROMPT_SUFFIXES[prompt_style] | |
| if suffix is None: | |
| return instr | |
| return f"{instr}\n{suffix}" | |
| # ── Preprocessing (matches training pipeline) ──────────────────────────────── | |
| def resize_to_target(image, target_size=TARGET_SIZE, interp=cv2.INTER_LINEAR): | |
| h, w = image.shape[:2] | |
| scale = target_size / max(h, w) | |
| new_w, new_h = int(round(w * scale)), int(round(h * scale)) | |
| resized = cv2.resize(image, (new_w, new_h), interpolation=interp) | |
| return resized, scale | |
| def prepare_inputs_from_arrays(rgb, depth, tmp_dir): | |
| rgb_resized, scale = resize_to_target(rgb, TARGET_SIZE, cv2.INTER_LINEAR) | |
| depth_resized, _ = resize_to_target(depth, TARGET_SIZE, cv2.INTER_NEAREST) | |
| new_h, new_w = rgb_resized.shape[:2] | |
| mask = ((depth_resized > 0.01) & (depth_resized < 10.0)).astype(np.float32) | |
| pcd = np.zeros((new_h, new_w, 4), dtype=np.float32) | |
| pcd[:, :, 2] = depth_resized | |
| pcd[:, :, 3] = mask | |
| img_path = os.path.join(tmp_dir, "frame_0_640.png") | |
| npz_path = os.path.join(tmp_dir, "frame_0_640.npz") | |
| cv2.imwrite(img_path, cv2.cvtColor(rgb_resized, cv2.COLOR_RGB2BGR)) | |
| np.savez_compressed(npz_path, pcd=pcd.astype(np.float16)) | |
| return img_path, npz_path, rgb_resized, depth_resized, scale | |
| def prepare_inputs(rgb_pil, depth_npy_path, tmp_dir): | |
| rgb = np.array(rgb_pil.convert("RGB")) | |
| depth = np.load(depth_npy_path).astype(np.float32) | |
| img_path, npz_path, rgb_resized, depth_resized, _ = prepare_inputs_from_arrays( | |
| rgb, depth, tmp_dir | |
| ) | |
| return img_path, npz_path, rgb_resized, depth_resized | |
| # ── Parsing ────────────────────────────────────────────────────────────────── | |
| def parse_trajectory(output): | |
| """Legacy <ans>...</ans> parser.""" | |
| waypoints, actions = [], [] | |
| ans_match = re.search(r"<ans>(.*?)</ans>", output, re.DOTALL) | |
| if not ans_match: | |
| ans_match = re.search(r"\[\[.*?\]\]", output, re.DOTALL) | |
| if not ans_match: | |
| return [], [] | |
| content = ans_match.group(0) | |
| else: | |
| content = ans_match.group(1) | |
| coord_pat = r"\[(\d+(?:\.\d+)?),\s*(\d+(?:\.\d+)?),\s*(\d+(?:\.\d+)?)\]" | |
| action_pat = r"<action>(.*?)</action>" | |
| parts = re.split(action_pat, content) | |
| for i, part in enumerate(parts): | |
| if i % 2 == 0: | |
| for c in re.findall(coord_pat, part): | |
| waypoints.append([float(c[0]), float(c[1]), float(c[2])]) | |
| actions.append(None) | |
| else: | |
| if actions: | |
| actions[-1] = part.strip() | |
| return waypoints, actions | |
| def parse_gt_structured_json(gpt_value): | |
| m = re.search(r"```json\s*(.*?)\s*```", gpt_value, re.DOTALL) | |
| raw = m.group(1) if m else gpt_value.strip() | |
| if not raw.lstrip().startswith("["): | |
| arr = re.search(r"\[.*\]", raw, re.DOTALL) | |
| if not arr: | |
| return [], [] | |
| raw = arr.group(0) | |
| try: | |
| entries = json.loads(raw) | |
| except json.JSONDecodeError: | |
| return [], [] | |
| if not isinstance(entries, list): | |
| return [], [] | |
| key = "point_3d" if any( | |
| isinstance(e, dict) and "point_3d" in e for e in entries | |
| ) else "point_2d" | |
| waypoints, actions = [], [] | |
| for entry in entries: | |
| if not isinstance(entry, dict) or key not in entry: | |
| continue | |
| pt = entry[key] | |
| waypoints.append( | |
| [float(pt[0]), float(pt[1]), float(pt[2]) if len(pt) > 2 else 0.0] | |
| ) | |
| grip = entry.get("gripper", "none") | |
| if grip == "close": | |
| actions.append("Close Gripper") | |
| elif grip == "open": | |
| actions.append("Open Gripper") | |
| else: | |
| actions.append(None) | |
| return waypoints, actions | |
| def parse_bbox_2d(output): | |
| m = re.search(r"```json\s*(.*?)\s*```", output, re.DOTALL) | |
| raw = m.group(1) if m else output.strip() | |
| if not raw.lstrip().startswith("["): | |
| arr = re.search(r"\[.*\]", raw, re.DOTALL) | |
| if not arr: | |
| return [] | |
| raw = arr.group(0) | |
| try: | |
| entries = json.loads(raw) | |
| except json.JSONDecodeError: | |
| return [] | |
| if not isinstance(entries, list): | |
| return [] | |
| boxes = [] | |
| for e in entries: | |
| if not isinstance(e, dict) or "bbox_2d" not in e: | |
| continue | |
| b = e["bbox_2d"] | |
| if len(b) < 4: | |
| continue | |
| boxes.append( | |
| (float(b[0]), float(b[1]), float(b[2]), float(b[3]), str(e.get("label", ""))) | |
| ) | |
| return boxes | |
| # ── 2D visualization ───────────────────────────────────────────────────────── | |
| COLOR_WP = (0, 255, 0) | |
| COLOR_GRASP = (255, 0, 0) | |
| COLOR_RELEASE = (0, 0, 255) | |
| COLOR_LINE = (255, 255, 0) | |
| def visualize_2d(image, waypoints, actions): | |
| if not waypoints: | |
| return image | |
| img = image.copy() | |
| h, w = img.shape[:2] | |
| pixels = [(int(wp[0] / 1000 * w), int(wp[1] / 1000 * h)) for wp in waypoints] | |
| for i in range(len(pixels) - 1): | |
| cv2.line(img, pixels[i], pixels[i + 1], COLOR_LINE, 2, cv2.LINE_AA) | |
| for i, (px, py) in enumerate(pixels): | |
| act = actions[i] if i < len(actions) else None | |
| if act and "Close" in act: | |
| color, r = COLOR_GRASP, 12 | |
| elif act and "Open" in act: | |
| color, r = COLOR_RELEASE, 12 | |
| else: | |
| color, r = COLOR_WP, 8 | |
| cv2.circle(img, (px, py), r, color, -1) | |
| cv2.circle(img, (px, py), r, (255, 255, 255), 2) | |
| cv2.putText(img, str(i), (px - 5, py + 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 2) | |
| cv2.putText( | |
| img, f"d={waypoints[i][2]:.2f}m", (px + 15, py), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1, | |
| ) | |
| y = 30 | |
| for label, color, xo in [ | |
| ("Waypoint", COLOR_WP, 10), ("Grasp", COLOR_GRASP, 110), ("Release", COLOR_RELEASE, 180) | |
| ]: | |
| cv2.circle(img, (xo, y - 5), 6, color, -1) | |
| cv2.putText(img, label, (xo + 10, y), cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1) | |
| return img | |
| def visualize_points(image, points): | |
| if not points: | |
| return image | |
| img = image.copy() | |
| h, w = img.shape[:2] | |
| for i, p in enumerate(points): | |
| px, py = int(p[0] / 1000 * w), int(p[1] / 1000 * h) | |
| cv2.circle(img, (px, py), 8, (0, 255, 0), -1) | |
| cv2.circle(img, (px, py), 8, (255, 255, 255), 2) | |
| cv2.putText( | |
| img, str(i + 1), (px + 11, py + 4), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, | |
| ) | |
| return img | |
| def visualize_bbox(image, boxes): | |
| if not boxes: | |
| return image | |
| img = image.copy() | |
| h, w = img.shape[:2] | |
| palette = [(0, 255, 0), (255, 80, 0), (0, 160, 255), (255, 0, 200), (255, 220, 0)] | |
| for i, (x1, y1, x2, y2, label) in enumerate(boxes): | |
| p1 = (int(x1 / 1000 * w), int(y1 / 1000 * h)) | |
| p2 = (int(x2 / 1000 * w), int(y2 / 1000 * h)) | |
| color = palette[i % len(palette)] | |
| cv2.rectangle(img, p1, p2, color, 2) | |
| tag = label or f"obj{i}" | |
| (tw, th), _ = cv2.getTextSize(tag, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1) | |
| cv2.rectangle(img, (p1[0], p1[1] - th - 6), (p1[0] + tw + 4, p1[1]), color, -1) | |
| cv2.putText(img, tag, (p1[0] + 2, p1[1] - 4), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1) | |
| return img | |
| # ── 3D scene + trajectory visualization ────────────────────────────────────── | |
| def colorize_depth(depth, min_d=0.1, max_d=3.0): | |
| d = np.clip(depth, min_d, max_d) | |
| d = ((d - min_d) / (max_d - min_d) * 255).astype(np.uint8) | |
| colored = cv2.applyColorMap(d, cv2.COLORMAP_TURBO) | |
| colored = cv2.cvtColor(colored, cv2.COLOR_BGR2RGB) | |
| colored[depth <= 0] = [0, 0, 0] | |
| return colored | |
| def _create_sphere_points(center, radius=0.005, n=150): | |
| phi = np.random.uniform(0, 2 * np.pi, n) | |
| ct = np.random.uniform(-1, 1, n) | |
| theta = np.arccos(ct) | |
| return np.stack([ | |
| center[0] + radius * np.sin(theta) * np.cos(phi), | |
| center[1] + radius * np.sin(theta) * np.sin(phi), | |
| center[2] + radius * np.cos(theta), | |
| ], axis=1) | |
| def _create_tube_points(p1, p2, radius=0.002): | |
| d = p2 - p1 | |
| L = np.linalg.norm(d) | |
| if L < 1e-8: | |
| return np.empty((0, 3)) | |
| d = d / L | |
| perp1 = np.cross(d, [1, 0, 0]) if abs(d[0]) < 0.9 else np.cross(d, [0, 1, 0]) | |
| perp1 /= np.linalg.norm(perp1) | |
| pts = [] | |
| for ti in np.linspace(0, 1, max(int(L / 0.002), 10)): | |
| c = p1 + ti * (p2 - p1) | |
| for a in np.linspace(0, 2 * np.pi, 6, endpoint=False): | |
| pts.append(c + radius * (np.cos(a) * perp1 + np.sin(a) * np.cross(d, perp1))) | |
| return np.array(pts) | |
| def _uvd_to_xyz(coords, intrinsics_3x3, img_w, img_h, uvd_norm=1000.0): | |
| K = np.array(intrinsics_3x3, dtype=np.float64) | |
| Kinv = np.linalg.inv(K) | |
| u_px = (coords[:, 0] / uvd_norm) * img_w | |
| v_px = (coords[:, 1] / uvd_norm) * img_h | |
| pixels = np.stack([u_px, v_px, np.ones(len(u_px))], axis=1) | |
| return coords[:, 2:3] * (pixels @ Kinv.T) | |
| def default_intrinsics(img_h, img_w): | |
| f = float(max(img_h, img_w)) | |
| return [[f, 0.0, img_w / 2.0], [0.0, f, img_h / 2.0], [0.0, 0.0, 1.0]] | |
| def build_scene_pcd_simple(rgb, depth, intrinsics_3x3, depth_trunc=3.0, stride=2): | |
| H, W = depth.shape[:2] | |
| if rgb.shape[:2] != (H, W): | |
| rgb = cv2.resize(rgb, (W, H)) | |
| K = np.asarray(intrinsics_3x3, dtype=np.float64) | |
| fx, fy, cx, cy = K[0, 0], K[1, 1], K[0, 2], K[1, 2] | |
| vs, us = np.mgrid[0:H:stride, 0:W:stride] | |
| z = depth[vs, us].astype(np.float32) | |
| valid = (z > 0.1) & (z < depth_trunc) | |
| z, us, vs = z[valid], us[valid], vs[valid] | |
| if z.size == 0: | |
| return None, None | |
| pts = np.stack([(us - cx) * z / fx, (vs - cy) * z / fy, z], axis=1) | |
| cols = rgb[vs, us].astype(np.float32) / 255.0 | |
| return pts, cols | |
| def build_3d_scene_figure(scene_pts, scene_cols, traj_dict, title="", subsample=2): | |
| fig = go.Figure() | |
| if scene_pts is not None and len(scene_pts) > 0: | |
| sp = scene_pts[::subsample] | |
| sr = (scene_cols[::subsample] * 255).astype(np.uint8) | |
| fig.add_trace(go.Scatter3d( | |
| x=sp[:, 0], y=sp[:, 1], z=sp[:, 2], mode="markers", | |
| marker=dict(size=1.5, color=[f"rgb({r},{g},{b})" for r, g, b in sr], opacity=0.6), | |
| name="Scene", hoverinfo="skip", | |
| )) | |
| for label, (xyz, color) in traj_dict.items(): | |
| if xyz is None or len(xyz) < 2: | |
| continue | |
| c = np.array(color) | |
| tpts, tcols = [], [] | |
| sp = _create_sphere_points(xyz[0], radius=0.008, n=400) | |
| tpts.append(sp) | |
| tcols.append(np.tile(np.clip(c * 1.3, 0, 1), (len(sp), 1))) | |
| ep = _create_sphere_points(xyz[-1], radius=0.008, n=400) | |
| tpts.append(ep) | |
| tcols.append(np.tile(c * 0.7, (len(ep), 1))) | |
| for i in range(len(xyz) - 1): | |
| tube = _create_tube_points(xyz[i], xyz[i + 1], radius=0.003) | |
| if len(tube) > 0: | |
| tpts.append(tube) | |
| tcols.append(np.tile(c, (len(tube), 1))) | |
| tpts = np.vstack(tpts) | |
| tcols = np.vstack(tcols) | |
| tr_rgb = (np.clip(tcols, 0, 1) * 255).astype(np.uint8) | |
| fig.add_trace(go.Scatter3d( | |
| x=tpts[:, 0], y=tpts[:, 1], z=tpts[:, 2], mode="markers", | |
| marker=dict(size=2.5, color=[f"rgb({r},{g},{b})" for r, g, b in tr_rgb], opacity=1.0), | |
| name=label, hoverinfo="skip", | |
| )) | |
| fig.update_layout( | |
| title=title, height=600, | |
| scene=dict( | |
| xaxis_title="X", yaxis_title="Y", zaxis_title="Z", | |
| aspectmode="data", bgcolor="white", | |
| camera=dict(eye=dict(x=0, y=0, z=-1.5), up=dict(x=0, y=-1, z=0)), | |
| ), | |
| paper_bgcolor="white", | |
| legend=dict(x=0.01, y=0.99, bgcolor="rgba(255,255,255,0.8)"), | |
| ) | |
| return fig | |
| def format_conversation(system, user, assistant): | |
| return ( | |
| "════════ SYSTEM ════════\n" | |
| f"{system.strip()}\n\n" | |
| "════════ USER ════════\n" | |
| f"{user.strip()}\n\n" | |
| "════════ ASSISTANT ════════\n" | |
| f"{assistant.strip()}" | |
| ) | |
| # ── Model (loaded once at module scope, eager .to("cuda")) ─────────────────── | |
| print(f"Downloading checkpoint {MODEL_ID} …") | |
| CKPT_DIR = snapshot_download(MODEL_ID) | |
| print(f"Checkpoint at {CKPT_DIR}") | |
| # Register the custom Qwen3-VL geometry model class with transformers Auto* . | |
| from hamster3d.model import register_qwen3_vl_geometry # noqa: E402 | |
| try: | |
| register_qwen3_vl_geometry() | |
| except Exception as e: # pragma: no cover | |
| print(f"register_qwen3_vl_geometry: {e!r}") | |
| from transformers import AutoModelForImageTextToText, AutoProcessor # noqa: E402 | |
| print("Loading processor …") | |
| PROCESSOR = AutoProcessor.from_pretrained(CKPT_DIR, trust_remote_code=True) | |
| print("Loading model …") | |
| MODEL = AutoModelForImageTextToText.from_pretrained( | |
| CKPT_DIR, | |
| dtype=torch.bfloat16, | |
| trust_remote_code=True, | |
| ).to("cuda") | |
| MODEL.eval() | |
| _PARAM_DTYPE = next(MODEL.parameters()).dtype | |
| print(f"Model loaded ({sum(p.numel() for p in MODEL.parameters()):,} params, dtype={_PARAM_DTYPE})") | |
| def _run_model(image_path, npz_path, query, system_prompt=V5_SYSTEM_PROMPT): | |
| """Greedy generation matching the reference ModelServer.predict.""" | |
| device = "cuda" | |
| rgb = np.array(Image.open(image_path).convert("RGB")) | |
| pcd = np.load(npz_path)["pcd"] # (H, W, 4) float16 | |
| depth = pcd[:, :, 2].astype(np.float32) | |
| rgb_tensor = torch.from_numpy(rgb).float().permute(2, 0, 1).unsqueeze(0) / 255.0 | |
| depth_tensor = torch.from_numpy(depth).float().unsqueeze(0) | |
| geometry_encoder_inputs = [rgb_tensor.to(device=device, dtype=_PARAM_DTYPE)] | |
| depth_maps = [depth_tensor.to(device=device, dtype=_PARAM_DTYPE)] | |
| messages = [] | |
| if system_prompt: | |
| messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]}) | |
| messages.append({ | |
| "role": "user", | |
| "content": [{"type": "image"}, {"type": "text", "text": query}], | |
| }) | |
| text = PROCESSOR.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) | |
| model_inputs = PROCESSOR( | |
| text=[text], | |
| images=[Image.open(image_path).convert("RGB")], | |
| return_tensors="pt", | |
| ).to(device) | |
| for _k, _v in list(model_inputs.items()): | |
| if torch.is_tensor(_v) and torch.is_floating_point(_v): | |
| model_inputs[_k] = _v.to(_PARAM_DTYPE) | |
| model_inputs["geometry_encoder_inputs"] = geometry_encoder_inputs | |
| model_inputs["depth_maps"] = depth_maps | |
| with torch.inference_mode(): | |
| output_ids = MODEL.generate( | |
| **model_inputs, | |
| max_new_tokens=1024, | |
| do_sample=False, | |
| temperature=None, | |
| top_p=None, | |
| ) | |
| input_len = model_inputs["input_ids"].shape[1] | |
| generated_ids = output_ids[:, input_len:] | |
| return PROCESSOR.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
| # ── Inference handler (ZeroGPU) ────────────────────────────────────────────── | |
| def predict(rgb_image, depth_file, instruction, prompt_style): | |
| """Predict a robot manipulation trajectory / pointing / bbox / VQA answer. | |
| Args: | |
| rgb_image: RGB scene image (PIL). Auto-resized to longest edge 640. | |
| depth_file: metric depth map as a .npy file (float32, meters, aligned to RGB). | |
| instruction: free-form language instruction (e.g. "Pick up the red cup"). | |
| prompt_style: one of "3D Trajectory", "2D Trajectory", "3D Pointing", | |
| "2D Pointing", "2D Bounding Box", "General VQA". | |
| Returns: | |
| (overlay_image, raw_output_text, conversation_text, plotly_3d_figure) | |
| """ | |
| if rgb_image is None: | |
| return None, "Please provide an RGB image.", "", None | |
| if depth_file is None: | |
| return None, "Please provide a metric depth .npy file.", "", None | |
| if not instruction or not instruction.strip(): | |
| return None, "Please enter a task instruction.", "", None | |
| depth_path = depth_file if isinstance(depth_file, str) else depth_file.name | |
| tmp_dir = tempfile.mkdtemp(prefix="hamster3d_") | |
| img_path, npz_path, rgb_resized, depth_resized = prepare_inputs( | |
| rgb_image, depth_path, tmp_dir | |
| ) | |
| h, w = rgb_resized.shape[:2] | |
| human_msg = build_v5_human_message(instruction, prompt_style) | |
| raw = _run_model(img_path, npz_path, human_msg, system_prompt=V5_SYSTEM_PROMPT) | |
| conversation = format_conversation(V5_SYSTEM_PROMPT, f"<image>{human_msg}", raw) | |
| K = default_intrinsics(h, w) | |
| scene_pts, scene_cols = build_scene_pcd_simple(rgb_resized, depth_resized, K) | |
| # General VQA → free-form answer, scene with no trajectory. | |
| if prompt_style == VQA_STYLE: | |
| fig = build_3d_scene_figure(scene_pts, scene_cols, {}, title=instruction.strip()) | |
| return rgb_resized, raw, conversation, fig | |
| # 2D bounding box. | |
| if prompt_style == BBOX_STYLE: | |
| boxes = parse_bbox_2d(raw) | |
| viz = visualize_bbox(rgb_resized, boxes) | |
| fig = build_3d_scene_figure(scene_pts, scene_cols, {}, title=instruction.strip()) | |
| return viz, raw, conversation, fig | |
| # Pointing → independent numbered points. | |
| if prompt_style in POINTING_STYLES: | |
| pts, _ = parse_gt_structured_json(raw) | |
| viz = visualize_points(rgb_resized, pts) if pts else rgb_resized | |
| fig = build_3d_scene_figure(scene_pts, scene_cols, {}, title=instruction.strip()) | |
| if pts: | |
| arr = np.array(pts, dtype=np.float32) | |
| if prompt_style == "2D Pointing" and depth_resized is not None: | |
| for i in range(len(arr)): | |
| up = int(np.clip(round(arr[i, 0] / 1000 * w), 0, w - 1)) | |
| vp = int(np.clip(round(arr[i, 1] / 1000 * h), 0, h - 1)) | |
| arr[i, 2] = float(depth_resized[vp, up]) | |
| xyz = _uvd_to_xyz(arr, K, w, h) | |
| fig.add_trace(go.Scatter3d( | |
| x=xyz[:, 0], y=xyz[:, 1], z=xyz[:, 2], mode="markers+text", | |
| marker=dict(size=8, color="lime", line=dict(width=2, color="white")), | |
| text=[str(i + 1) for i in range(len(xyz))], textposition="top center", | |
| name="Points", hoverinfo="skip", | |
| )) | |
| return viz, raw, conversation, fig | |
| # Trajectory (2D / 3D). | |
| waypoints, actions = parse_gt_structured_json(raw) | |
| if not waypoints: | |
| waypoints, actions = parse_trajectory(raw) | |
| viz_2d = visualize_2d(rgb_resized, waypoints, actions) if waypoints else rgb_resized | |
| traj_dict = {} | |
| if waypoints: | |
| pred_xyz = _uvd_to_xyz(np.array(waypoints, dtype=np.float32), K, w, h) | |
| traj_dict["Predicted trajectory"] = (pred_xyz, [1.0, 0.2, 0.0]) | |
| fig = build_3d_scene_figure(scene_pts, scene_cols, traj_dict, title=instruction.strip()) | |
| return viz_2d, raw, conversation, fig | |
| # ── Examples browser helpers ───────────────────────────────────────────────── | |
| def _load_example_assets(idx): | |
| """Return (rgb_pil, depth_npy_path, instruction) for bundled example idx.""" | |
| prefix = f"sample_{int(idx)}" | |
| rgb_path = os.path.join(EXAMPLES_DIR, f"{prefix}_rgb.png") | |
| depth_path = os.path.join(EXAMPLES_DIR, f"{prefix}_depth.npy") | |
| instr_path = os.path.join(EXAMPLES_DIR, f"{prefix}_instruction.txt") | |
| instruction = "" | |
| if os.path.isfile(instr_path): | |
| instruction = open(instr_path).read().strip() | |
| return rgb_path, depth_path, instruction | |
| def _example_rows(): | |
| rows = [] | |
| for i in range(6): | |
| rgb_path, depth_path, instruction = _load_example_assets(i) | |
| if os.path.isfile(rgb_path) and os.path.isfile(depth_path): | |
| rows.append([rgb_path, depth_path, instruction, V5_DEFAULT_STYLE]) | |
| return rows | |
| def on_depth_upload(depth_file): | |
| if depth_file is None: | |
| return None | |
| try: | |
| path = depth_file if isinstance(depth_file, str) else depth_file.name | |
| depth = np.load(path).astype(np.float32) | |
| return colorize_depth(depth) | |
| except Exception: | |
| return None | |
| # ── UI ─────────────────────────────────────────────────────────────────────── | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| .mono textarea { font-family: monospace; font-size: 13px; } | |
| """ | |
| with gr.Blocks(title="3D HAMSTER") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# 🐹 3D HAMSTER — 3D Trajectory Prediction\n" | |
| "Depth-aware **Vision-Language-Action** planner (Qwen3-VL-8B + LingBot-Depth " | |
| "geometry encoder). From a single **RGB image + metric depth map + language " | |
| "instruction**, it predicts a metric **3D end-effector trajectory** " | |
| "(`[u, v, depth]` waypoints + gripper states), plus 2D trajectory / pointing / " | |
| "bounding-box / VQA modes.\n\n" | |
| "[Paper](https://huggingface.co/papers/2606.31329) · " | |
| "[Model](https://huggingface.co/DAVIAN-Robotics/3D_HAMSTER) · " | |
| "[Code](https://github.com/DAVIAN-Robotics/3D_HAMSTER) · " | |
| "[Project page](https://davian-robotics.github.io/3D_HAMSTER/)\n\n" | |
| "> ⚠️ Depth must be **metric (meters)** and aligned to the RGB frame. " | |
| "Disparity / normalized / millimeter depth will degrade the geometry." | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| rgb_input = gr.Image(label="RGB Image", type="pil", height=320) | |
| depth_input = gr.File( | |
| label="Metric Depth (.npy, float32, meters)", file_types=[".npy"] | |
| ) | |
| depth_preview = gr.Image(label="Depth preview", type="numpy", height=180) | |
| instruction_input = gr.Textbox( | |
| label="Task instruction", | |
| placeholder="e.g. Pick up the red block and place it on the blue plate.", | |
| lines=2, | |
| ) | |
| prompt_style_input = gr.Radio( | |
| label="Prompt style", | |
| choices=V5_PROMPT_STYLES, | |
| value=V5_DEFAULT_STYLE, | |
| info="Trajectory: manipulation waypoints | Pointing: object points | " | |
| "2D Bounding Box: object box | General VQA: free-form answer", | |
| ) | |
| run_btn = gr.Button("Predict trajectory", variant="primary", size="lg") | |
| with gr.Column(scale=1): | |
| overlay_output = gr.Image(label="2D overlay", type="numpy", height=340) | |
| plot_3d_output = gr.Plot(label="3D scene + trajectory (rotate / zoom)") | |
| with gr.Accordion("Model output", open=False): | |
| raw_output = gr.Textbox( | |
| label="Raw model output", lines=4, interactive=False, | |
| elem_classes=["mono"], | |
| ) | |
| conversation_output = gr.Textbox( | |
| label="Full conversation", lines=10, interactive=False, | |
| elem_classes=["mono"], | |
| ) | |
| gr.Examples( | |
| examples=_example_rows(), | |
| inputs=[rgb_input, depth_input, instruction_input, prompt_style_input], | |
| outputs=[overlay_output, raw_output, conversation_output, plot_3d_output], | |
| fn=predict, | |
| cache_examples=False, | |
| run_on_click=True, | |
| label="Bundled examples (RGB + depth + instruction)", | |
| ) | |
| depth_input.change(fn=on_depth_upload, inputs=[depth_input], outputs=[depth_preview]) | |
| run_btn.click( | |
| fn=predict, | |
| inputs=[rgb_input, depth_input, instruction_input, prompt_style_input], | |
| outputs=[overlay_output, raw_output, conversation_output, plot_3d_output], | |
| api_name="predict", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |