app.py

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
# ---------------------------------------------------------------------------
@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)