Spaces:
Running on Zero
Running on Zero
| import os | |
| import sys | |
| import spaces # must be imported before torch | |
| import tempfile | |
| import time | |
| from collections import OrderedDict | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| import trimesh | |
| from huggingface_hub import hf_hub_download | |
| from plyfile import PlyData | |
| from scipy.spatial import cKDTree | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| from litept.model import LitePT # noqa: E402 | |
| REPO_ID = "prs-eth/LitePT" | |
| HERE = os.path.dirname(os.path.abspath(__file__)) | |
| # --------------------------------------------------------------------------- | |
| # Class names / palettes | |
| # --------------------------------------------------------------------------- | |
| NUSCENES_NAMES = [ | |
| "barrier", "bicycle", "bus", "car", "construction_vehicle", "motorcycle", | |
| "pedestrian", "traffic_cone", "trailer", "truck", "driveable_surface", | |
| "other_flat", "sidewalk", "terrain", "manmade", "vegetation", | |
| ] | |
| NUSCENES_COLORS = [ | |
| (255, 120, 50), (255, 192, 203), (255, 255, 0), (0, 150, 245), (0, 255, 255), | |
| (200, 180, 0), (255, 0, 0), (255, 240, 150), (135, 60, 0), (160, 32, 240), | |
| (255, 0, 255), (139, 137, 137), (75, 0, 75), (150, 240, 80), (230, 230, 250), | |
| (0, 175, 0), | |
| ] | |
| SCANNET_NAMES = [ | |
| "wall", "floor", "cabinet", "bed", "chair", "sofa", "table", "door", "window", | |
| "bookshelf", "picture", "counter", "desk", "curtain", "refrigerator", | |
| "shower curtain", "toilet", "sink", "bathtub", "otherfurniture", | |
| ] | |
| SCANNET_COLORS = [ | |
| (174, 199, 232), (152, 223, 138), (31, 119, 180), (255, 187, 120), (188, 189, 34), | |
| (140, 86, 75), (255, 152, 150), (214, 39, 40), (197, 176, 213), (148, 103, 189), | |
| (196, 156, 148), (23, 190, 207), (247, 182, 210), (219, 219, 141), (255, 127, 14), | |
| (158, 218, 229), (44, 160, 44), (112, 128, 144), (227, 119, 194), (82, 84, 163), | |
| ] | |
| S3D_NAMES = [ | |
| "wall", "floor", "cabinet", "bed", "chair", "sofa", "table", "door", "window", | |
| "picture", "desk", "shelves", "curtain", "dresser", "pillow", "mirror", "ceiling", | |
| "refrigerator", "television", "nightstand", "sink", "lamp", "otherstructure", | |
| "otherfurniture", "otherprop", | |
| ] | |
| S3D_COLORS = [ | |
| (174, 199, 232), (152, 223, 138), (31, 119, 180), (255, 187, 120), (188, 189, 34), | |
| (140, 86, 75), (255, 152, 150), (214, 39, 40), (197, 176, 213), (196, 156, 148), | |
| (247, 182, 210), (148, 103, 189), (219, 219, 141), (178, 76, 76), (255, 215, 0), | |
| (0, 191, 255), (230, 230, 230), (255, 127, 14), (60, 60, 60), (205, 133, 63), | |
| (112, 128, 144), (255, 255, 102), (120, 120, 120), (82, 84, 163), (100, 200, 200), | |
| ] | |
| SMALL = dict( | |
| enc_depths=(2, 2, 2, 6, 2), enc_channels=(36, 72, 144, 252, 504), | |
| enc_num_head=(2, 4, 8, 14, 28), dec_channels=(72, 72, 144, 252), | |
| dec_num_head=(4, 4, 8, 14), | |
| ) | |
| LARGE = dict( | |
| enc_depths=(3, 3, 3, 12, 3), enc_channels=(72, 144, 288, 576, 864), | |
| enc_num_head=(4, 8, 16, 32, 48), dec_channels=(72, 144, 288, 576), | |
| dec_num_head=(4, 8, 16, 32), | |
| ) | |
| MODELS = OrderedDict( | |
| [ | |
| ( | |
| "Outdoor LiDAR · nuScenes (LitePT-S)", | |
| dict(ckpt="nuscenes-semseg-litept-small-v1m1/model/model_best.pth", | |
| arch=SMALL, in_channels=4, grid=0.05, kind="outdoor", | |
| names=NUSCENES_NAMES, colors=NUSCENES_COLORS), | |
| ), | |
| ( | |
| "Indoor RGB · ScanNet (LitePT-S)", | |
| dict(ckpt="scannet-semseg-litept-small-v1m1/model/model_best.pth", | |
| arch=SMALL, in_channels=6, grid=0.02, kind="indoor", | |
| names=SCANNET_NAMES, colors=SCANNET_COLORS), | |
| ), | |
| ( | |
| "Indoor RGB · Structured3D, synthetic (LitePT-L)", | |
| dict(ckpt="structured3d-semseg-litept-large-v1m1/model/model_best.pth", | |
| arch=LARGE, in_channels=6, grid=0.02, kind="indoor", | |
| names=S3D_NAMES, colors=S3D_COLORS), | |
| ), | |
| ] | |
| ) | |
| MODEL_CHOICES = list(MODELS.keys()) | |
| OUTDOOR_DEFAULT = MODEL_CHOICES[0] | |
| INDOOR_DEFAULT = MODEL_CHOICES[1] | |
| def build_model(cfg): | |
| backbone = LitePT(in_channels=cfg["in_channels"], **cfg["arch"]) | |
| head = torch.nn.Linear(72, len(cfg["names"])) | |
| path = hf_hub_download(REPO_ID, cfg["ckpt"]) | |
| sd = torch.load(path, map_location="cpu", weights_only=False)["state_dict"] | |
| bb, hd = OrderedDict(), OrderedDict() | |
| for k, v in sd.items(): | |
| k = k[len("module."):] if k.startswith("module.") else k | |
| if k.startswith("backbone."): | |
| bb[k[len("backbone."):]] = v | |
| elif k.startswith("seg_head."): | |
| hd[k[len("seg_head."):]] = v | |
| backbone.load_state_dict(bb, strict=True) | |
| head.load_state_dict(hd, strict=True) | |
| del sd | |
| return backbone.eval().to("cuda"), head.eval().to("cuda") | |
| NETS = {name: build_model(cfg) for name, cfg in MODELS.items()} | |
| print("All LitePT checkpoints loaded.") | |
| # --------------------------------------------------------------------------- | |
| # I/O helpers | |
| # --------------------------------------------------------------------------- | |
| MAX_POINTS = 2_000_000 | |
| def estimate_normals(xyz, k=16): | |
| """PCA normals from k nearest neighbours (used when the input has none).""" | |
| tree = cKDTree(xyz) | |
| _, idx = tree.query(xyz, k=min(k, len(xyz))) | |
| nb = xyz[idx] - xyz[idx].mean(1, keepdims=True) | |
| cov = np.einsum("nki,nkj->nij", nb, nb) | |
| _, vecs = np.linalg.eigh(cov) | |
| n = vecs[:, :, 0] | |
| # orient towards the cloud's upper centre (rough, but consistent) | |
| view = xyz.mean(0) + np.array([0, 0, 2.0]) | |
| flip = ((view - xyz) * n).sum(1) < 0 | |
| n[flip] *= -1 | |
| return n.astype(np.float32) | |
| def load_point_cloud(path): | |
| """Returns dict(coord[N,3], color[N,3] in 0..255 or None, normal or None, strength[N,1] or None).""" | |
| ext = os.path.splitext(path)[1].lower() | |
| color = normal = strength = None | |
| if ext == ".ply": | |
| v = PlyData.read(path)["vertex"].data | |
| names = v.dtype.names | |
| coord = np.stack([v["x"], v["y"], v["z"]], 1).astype(np.float32) | |
| if all(c in names for c in ("red", "green", "blue")): | |
| color = np.stack([v["red"], v["green"], v["blue"]], 1).astype(np.float32) | |
| if color.max() <= 1.0: | |
| color = color * 255.0 | |
| if all(c in names for c in ("nx", "ny", "nz")): | |
| normal = np.stack([v["nx"], v["ny"], v["nz"]], 1).astype(np.float32) | |
| for key in ("intensity", "strength", "scalar_intensity", "reflectance"): | |
| if key in names: | |
| strength = np.asarray(v[key], np.float32).reshape(-1, 1) | |
| break | |
| elif ext == ".bin": | |
| raw = np.fromfile(path, dtype=np.float32) | |
| cols = 5 if raw.size % 5 == 0 else 4 # nuScenes (x,y,z,i,ring) or KITTI (x,y,z,i) | |
| if raw.size % cols != 0: | |
| raise gr.Error("Could not parse .bin: expected float32 N×5 (nuScenes) or N×4 (KITTI).") | |
| raw = raw.reshape(-1, cols) | |
| coord, strength = raw[:, :3], raw[:, 3:4] | |
| elif ext in (".npy", ".txt", ".xyz", ".pts", ".csv"): | |
| if ext == ".npy": | |
| arr = np.load(path) | |
| else: | |
| arr = np.loadtxt(path, delimiter="," if ext == ".csv" else None, comments=("#", "//")) | |
| arr = np.asarray(arr, np.float32) | |
| if arr.ndim != 2 or arr.shape[1] < 3: | |
| raise gr.Error("Array input must be N×C with C ≥ 3 (x y z [r g b] [nx ny nz] or x y z intensity).") | |
| coord = arr[:, :3] | |
| if arr.shape[1] == 4: | |
| strength = arr[:, 3:4] | |
| elif arr.shape[1] >= 6: | |
| color = arr[:, 3:6] | |
| if color.max() <= 1.0: | |
| color = color * 255.0 | |
| if arr.shape[1] >= 9: | |
| normal = arr[:, 6:9] | |
| else: | |
| raise gr.Error(f"Unsupported file type '{ext}'. Use .ply, .bin, .npy, .txt, .xyz or .csv.") | |
| coord = np.ascontiguousarray(coord, np.float32) | |
| ok = np.isfinite(coord).all(1) | |
| if not ok.all(): | |
| coord = coord[ok] | |
| color = color[ok] if color is not None else None | |
| normal = normal[ok] if normal is not None else None | |
| strength = strength[ok] if strength is not None else None | |
| if len(coord) < 16: | |
| raise gr.Error("The point cloud has too few points.") | |
| if len(coord) > MAX_POINTS: | |
| sel = np.random.default_rng(0).choice(len(coord), MAX_POINTS, replace=False) | |
| coord = coord[sel] | |
| color = color[sel] if color is not None else None | |
| normal = normal[sel] if normal is not None else None | |
| strength = strength[sel] if strength is not None else None | |
| return dict(coord=coord, color=color, normal=normal, strength=strength) | |
| def grid_sample(coord, grid_size, seed=0): | |
| """Pointcept GridSample (train mode): keep one point per voxel + inverse map.""" | |
| grid = np.floor(coord / grid_size).astype(np.int64) | |
| grid -= grid.min(0) | |
| _, idx_sort_unique, inverse, count = np.unique( | |
| grid, axis=0, return_index=True, return_inverse=True, return_counts=True | |
| ) | |
| inverse = inverse.reshape(-1) | |
| order = np.argsort(inverse, kind="stable") | |
| starts = np.concatenate([[0], np.cumsum(count)[:-1]]) | |
| rng = np.random.default_rng(seed) | |
| pick = order[starts + rng.integers(0, count.max(), count.size) % count] | |
| return pick, inverse, grid[pick] | |
| def prepare_input(pc, cfg): | |
| coord = pc["coord"].copy() | |
| if cfg["kind"] == "outdoor": | |
| strength = pc["strength"] | |
| if strength is None: | |
| strength = np.zeros((len(coord), 1), np.float32) | |
| elif strength.max() > 1.0: | |
| strength = strength / 255.0 # nuScenes convention | |
| pick, inverse, grid_coord = grid_sample(coord, cfg["grid"]) | |
| c = coord[pick] | |
| feat = np.concatenate([c, strength[pick]], 1) | |
| else: | |
| color = pc["color"] | |
| if color is None: | |
| color = np.full((len(coord), 3), 127.0, np.float32) | |
| normal = pc["normal"] | |
| if normal is None: | |
| normal = estimate_normals(coord) | |
| # CenterShift(apply_z=True) -> NormalizeColor -> GridSample -> CenterShift(apply_z=False) | |
| mn, mx = coord.min(0), coord.max(0) | |
| coord = coord - np.array([(mn[0] + mx[0]) / 2, (mn[1] + mx[1]) / 2, mn[2]], np.float32) | |
| color = color / 255.0 | |
| pick, inverse, grid_coord = grid_sample(coord, cfg["grid"]) | |
| c = coord[pick] | |
| cmn, cmx = c.min(0), c.max(0) | |
| c = c - np.array([(cmn[0] + cmx[0]) / 2, (cmn[1] + cmx[1]) / 2, 0], np.float32) | |
| feat = np.concatenate([color[pick], normal[pick]], 1) | |
| return dict( | |
| coord=torch.from_numpy(np.ascontiguousarray(c, np.float32)), | |
| grid_coord=torch.from_numpy(grid_coord.astype(np.int32)), | |
| feat=torch.from_numpy(np.ascontiguousarray(feat, np.float32)), | |
| offset=torch.tensor([len(c)], dtype=torch.long), | |
| grid_size=cfg["grid"], | |
| ), inverse | |
| def to_glb(coord, rgb, max_points=400_000): | |
| if len(coord) > max_points: | |
| sel = np.random.default_rng(0).choice(len(coord), max_points, replace=False) | |
| coord, rgb = coord[sel], rgb[sel] | |
| c = coord - coord.mean(0) | |
| # z-up (point cloud) -> y-up (glTF) | |
| c = np.stack([c[:, 0], c[:, 2], -c[:, 1]], 1) | |
| rgba = np.concatenate([rgb.astype(np.uint8), np.full((len(rgb), 1), 255, np.uint8)], 1) | |
| scene = trimesh.Scene(trimesh.PointCloud(vertices=c, colors=rgba)) | |
| f = tempfile.NamedTemporaryFile(suffix=".glb", delete=False) | |
| scene.export(f.name) | |
| return f.name | |
| def input_colors(pc): | |
| if pc["color"] is not None: | |
| return np.clip(pc["color"], 0, 255) | |
| if pc["strength"] is not None: | |
| s = pc["strength"][:, 0] | |
| s = (s - s.min()) / (np.ptp(s) + 1e-6) | |
| s = np.sqrt(s) | |
| else: | |
| z = pc["coord"][:, 2] | |
| s = (z - np.percentile(z, 1)) / (np.ptp(z) + 1e-6) | |
| s = np.clip(s, 0, 1) | |
| base = np.array([40, 60, 120], np.float32) | |
| top = np.array([250, 230, 120], np.float32) | |
| return base + s[:, None] * (top - base) | |
| def legend_html(labels, cfg): | |
| counts = np.bincount(labels, minlength=len(cfg["names"])) | |
| total = counts.sum() | |
| rows = [] | |
| for i in np.argsort(-counts): | |
| if counts[i] == 0: | |
| continue | |
| r, g, b = cfg["colors"][i] | |
| pct = 100.0 * counts[i] / total | |
| rows.append( | |
| f"<div style='display:flex;align-items:center;gap:8px;margin:2px 0'>" | |
| f"<span style='width:14px;height:14px;border-radius:3px;background:rgb({r},{g},{b});" | |
| f"display:inline-block;border:1px solid #8884'></span>" | |
| f"<span style='flex:1'>{cfg['names'][i]}</span>" | |
| f"<span style='opacity:.75;font-variant-numeric:tabular-nums'>{pct:.1f}%</span></div>" | |
| ) | |
| return "<div style='columns:2;column-gap:24px;font-size:14px'>" + "".join(rows) + "</div>" | |
| # --------------------------------------------------------------------------- | |
| # Inference | |
| # --------------------------------------------------------------------------- | |
| def _forward(model_name, data): | |
| backbone, head = NETS[model_name] | |
| with torch.no_grad(): | |
| inp = {k: (v.to("cuda") if isinstance(v, torch.Tensor) else v) for k, v in data.items()} | |
| point = backbone(inp) | |
| logits = head(point.feat) | |
| return logits.argmax(-1).cpu().numpy().astype(np.int64) | |
| def preview(file_path: str): | |
| """Render the uploaded point cloud (RGB, intensity or height colouring) as a GLB.""" | |
| if not file_path: | |
| return None | |
| pc = load_point_cloud(file_path) | |
| return to_glb(pc["coord"], input_colors(pc)) | |
| def segment(file_path: str, model_name: str = INDOOR_DEFAULT): | |
| """Semantic segmentation of a 3D point cloud with LitePT. | |
| Args: | |
| file_path: point cloud file (.ply with x/y/z [+ red/green/blue, nx/ny/nz, intensity], | |
| .bin nuScenes/KITTI LiDAR sweep, or .npy/.txt/.xyz/.csv N×C array). Z must be up, units in meters. | |
| model_name: which LitePT checkpoint to use (outdoor nuScenes, indoor ScanNet or Structured3D). | |
| Returns: | |
| input preview GLB, segmentation GLB, legend HTML, per-point label file (.npy), status text. | |
| """ | |
| if not file_path: | |
| raise gr.Error("Please upload a point cloud or pick an example.") | |
| if model_name not in MODELS: | |
| raise gr.Error(f"Unknown model '{model_name}'.") | |
| cfg = MODELS[model_name] | |
| t0 = time.perf_counter() | |
| pc = load_point_cloud(file_path) | |
| data, inverse = prepare_input(pc, cfg) | |
| t1 = time.perf_counter() | |
| labels_vox = _forward(model_name, data) | |
| t2 = time.perf_counter() | |
| labels = labels_vox[inverse] | |
| palette = np.array(cfg["colors"], np.uint8) | |
| seg_glb = to_glb(pc["coord"], palette[labels]) | |
| in_glb = to_glb(pc["coord"], input_colors(pc)) | |
| lab_file = tempfile.NamedTemporaryFile(suffix="_labels.npy", delete=False).name | |
| np.save(lab_file, labels.astype(np.int16)) | |
| status = ( | |
| f"{len(pc['coord']):,} points → {len(data['coord']):,} voxels " | |
| f"({cfg['grid'] * 100:.0f} cm grid) · GPU call {t2 - t1:.2f}s · preprocessing {t1 - t0:.2f}s" | |
| ) | |
| return in_glb, seg_glb, legend_html(labels, cfg), lab_file, status | |
| def model_for_file(file_path): | |
| if not file_path: | |
| return gr.update() | |
| if file_path.lower().endswith(".bin"): | |
| return gr.update(value=OUTDOOR_DEFAULT) | |
| return gr.update() | |
| # --------------------------------------------------------------------------- | |
| # UI | |
| # --------------------------------------------------------------------------- | |
| CSS = """ | |
| #col-container { max-width: 1200px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="LitePT") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# LitePT: Lighter Yet Stronger Point Transformer\n" | |
| "3D semantic segmentation of point clouds with **LitePT** (CVPR 2026). Upload an indoor RGB scan " | |
| "(`.ply`) or an outdoor LiDAR sweep (`.bin`), pick a checkpoint, and get per-point labels. \n" | |
| "[Paper](https://huggingface.co/papers/2512.13689) · [Code](https://github.com/prs-eth/LitePT) · " | |
| "[Weights](https://huggingface.co/prs-eth/LitePT) · [Project page](https://litept.github.io/)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| file_in = gr.File( | |
| label="Point cloud (.ply / .bin / .npy / .txt) — z-up, meters", | |
| file_types=[".ply", ".bin", ".npy", ".txt", ".xyz", ".csv", ".pts"], | |
| type="filepath", | |
| ) | |
| model_dd = gr.Dropdown(MODEL_CHOICES, value=INDOOR_DEFAULT, label="Checkpoint") | |
| run_btn = gr.Button("Segment", variant="primary") | |
| status = gr.Markdown() | |
| legend = gr.HTML(label="Classes") | |
| labels_out = gr.File(label="Per-point labels (.npy, int16, input order)") | |
| with gr.Column(scale=2): | |
| with gr.Row(): | |
| in_view = gr.Model3D(label="Input", height=480, clear_color=(0, 0, 0, 0)) | |
| seg_view = gr.Model3D(label="LitePT segmentation", height=480, clear_color=(0, 0, 0, 0)) | |
| with gr.Accordion("Input format notes", open=False): | |
| gr.Markdown( | |
| "- **Indoor models** (ScanNet / Structured3D) use RGB + normals, 2 cm voxels. Missing normals are " | |
| "estimated with PCA over 16 nearest neighbours; missing colours are set to grey.\n" | |
| "- **Outdoor model** (nuScenes) uses xyz + intensity in the sensor frame, 5 cm voxels. " | |
| "`.bin` files are read as nuScenes (N×5) or KITTI (N×4) float32 sweeps.\n" | |
| "- The scene must be **z-up** and in **meters**. Inputs above 2M points are randomly subsampled.\n" | |
| "- Single forward pass (no test-time augmentation), so scores can be slightly below the paper's numbers." | |
| ) | |
| gr.Examples( | |
| examples=[ | |
| ["examples/nuscenes_sample.bin", OUTDOOR_DEFAULT], | |
| ["examples/icl_nuim_livingroom.ply", INDOOR_DEFAULT], | |
| ], | |
| inputs=[file_in, model_dd], | |
| outputs=[in_view, seg_view, legend, labels_out, status], | |
| fn=segment, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| file_in.upload(model_for_file, inputs=file_in, outputs=model_dd, api_name=False).then( | |
| preview, inputs=file_in, outputs=in_view, api_name=False | |
| ) | |
| run_btn.click( | |
| segment, | |
| inputs=[file_in, model_dd], | |
| outputs=[in_view, seg_view, legend, labels_out, status], | |
| api_name="segment", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |