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"
" f"" f"{cfg['names'][i]}" f"{pct:.1f}%
" ) return "
" + "".join(rows) + "
" # --------------------------------------------------------------------------- # Inference # --------------------------------------------------------------------------- @spaces.GPU(duration=22) 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)