PatchAlign3D zero-shot 3D part segmentation demo

README.md

@@ -1,13 +1,50 @@
 ---
-title: Patch Align 3d
-emoji: 🦀
-colorFrom: indigo
-colorTo: purple
+title: PatchAlign3D
+emoji: 🧩
+colorFrom: gray
+colorTo: pink
 sdk: gradio
 sdk_version: 6.22.0
-python_version: '3.12'
 app_file: app.py
 pinned: false
+license: mit
+short_description: Zero-shot 3D part segmentation from text queries
+python_version: "3.12"
+startup_duration_timeout: 1h
+models:
+  - patchalign3d/patchalign3d-encoder
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PatchAlign3D · zero-shot 3D part segmentation
+
+Name the parts you want **in words** and see them highlighted on a 3D shape. A single forward pass of a
+point-cloud encoder whose *patch-level* features are aligned to CLIP text space — no test-time multi-view
+rendering, no per-category training.
+
+- Paper: https://huggingface.co/papers/2601.02457
+- Code: https://github.com/souhail-hadgi/PatchAlign3D
+- Weights: https://huggingface.co/patchalign3d/patchalign3d-encoder
+
+## How it works
+
+Points are centred and scaled to the unit sphere and the Y/Z axes are swapped to match the training
+convention (exactly as in the official `infer.py`). Furthest-point sampling picks patch centres, a k-NN
+neighbourhood around each becomes a patch token, and a 12-layer point transformer produces one feature per
+patch. A learned linear head projects those into the CLIP `ViT-bigG-14 (laion2b_s39b_b160k)` text space,
+where they are matched against a prompt ensemble. Each point takes the label of its nearest patch centre.
+
+## Implementation notes
+
+`patchalign3d.py` is a faithful port of the official inference path with three behaviour-preserving
+deviations, so the Space runs without custom CUDA extensions:
+
+1. `pointnet2_ops.furthest_point_sample` → pure-torch iterative FPS (same start index, squared distances,
+   argmax selection).
+2. `knn_cuda.KNN` → `torch.cdist` + `topk`.
+3. `open_clip ViT-bigG-14 / laion2b_s39b_b160k` text tower → the byte-identical (up to fp16 rounding)
+   `text_encoder_2` + `tokenizer_2` of `stabilityai/stable-diffusion-xl-base-1.0`, which avoids a 10 GB
+   download of the full image+text CLIP checkpoint.
+
+Verified against ShapeNetPart ground truth on CPU before deployment: the checkpoint loads with 0 missing
+and 0 unexpected keys, and per-shape zero-shot mIoU averages **0.59** with the default `part_only` prompts
+(**0.64** with `ensemble`) across 8 held-out test shapes.

app.py

@@ -0,0 +1,431 @@
+"""PatchAlign3D — open-vocabulary (zero-shot) 3D part segmentation from point clouds.
+
+Paper:   https://huggingface.co/papers/2601.02457
+Code:    https://github.com/souhail-hadgi/PatchAlign3D
+Weights: https://huggingface.co/patchalign3d/patchalign3d-encoder
+"""
+
+import os
+
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+import spaces  # noqa: E402  (must precede torch)
+
+import tempfile  # noqa: E402
+import time  # noqa: E402
+from pathlib import Path  # 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
+import trimesh  # noqa: E402
+from huggingface_hub import hf_hub_download  # noqa: E402
+from transformers import CLIPTextModelWithProjection, CLIPTokenizer  # noqa: E402
+
+import patchalign3d as pa  # noqa: E402
+
+# --------------------------------------------------------------------------------------
+# Models — module scope, eager .to("cuda"); ZeroGPU streams them in on the first call
+# --------------------------------------------------------------------------------------
+
+CKPT = hf_hub_download("patchalign3d/patchalign3d-encoder", "patchalign3d.pt")
+model, proj = pa.load_patchalign3d(CKPT)
+model = model.to("cuda")
+proj = proj.to("cuda")
+
+tokenizer = CLIPTokenizer.from_pretrained(pa.CLIP_TEXT_REPO, subfolder=pa.CLIP_TOKENIZER_SUBFOLDER)
+text_model = (
+    CLIPTextModelWithProjection.from_pretrained(
+        pa.CLIP_TEXT_REPO, subfolder=pa.CLIP_TEXT_SUBFOLDER, variant="fp16", dtype=torch.float32
+    )
+    .eval()
+    .to("cuda")
+)
+
+print(
+    f"[init] PatchAlign3D encoder {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M params | "
+    f"CLIP ViT-bigG-14 text tower {sum(p.numel() for p in text_model.parameters()) / 1e6:.1f}M params | "
+    f"tokenizer pad={tokenizer.pad_token_id} ctx={tokenizer.model_max_length}"
+)
+
+MAX_LABELS = 12
+MESH_EXTS = {".obj", ".glb", ".gltf", ".stl", ".off", ".ply", ".dae", ".3mf"}
+
+# Distinguishable qualitative palette
+PALETTE = [
+    "#e6194b", "#3cb44b", "#4363d8", "#f58231", "#911eb4", "#00b8d4",
+    "#f032e6", "#a1c800", "#fabed4", "#469990", "#9a6324", "#7f0000",
+]
+
+# ShapeNetPart part vocabularies, verbatim from the official eval.py
+PRESETS = {
+    "— custom —": ("", ""),
+    "Airplane": ("body, wing, tail, engine or frame", "airplane"),
+    "Bag": ("handle, body", "bag"),
+    "Cap": ("crown, brim", "cap"),
+    "Car": ("roof, hood, wheel, body", "car"),
+    "Chair": ("back, seat, leg, arm", "chair"),
+    "Earphone": ("earcup, headband, data wire", "earphone"),
+    "Guitar": ("headstock, neck, body", "guitar"),
+    "Knife": ("blade, handle", "knife"),
+    "Lamp": ("base, lampshade, fixing bracket, pole", "lamp"),
+    "Laptop": ("keyboard, screen", "laptop"),
+    "Motorbike": ("gas tank, seat, wheel, handles or handlebars, headlight, engine or frame", "motorbike"),
+    "Mug": ("handle, cup", "mug"),
+    "Pistol": ("barrel, handle or grip, trigger and guard", "pistol"),
+    "Rocket": ("body, fin, nose", "rocket"),
+    "Skateboard": ("wheel, deck, belt for foot", "skateboard"),
+    "Table": ("desktop, leg or support, drawer", "table"),
+}
+
+
+# --------------------------------------------------------------------------------------
+# Shape loading
+# --------------------------------------------------------------------------------------
+
+
+def _resample(pts: np.ndarray, npoints: int, seed: int) -> np.ndarray:
+    n = len(pts)
+    if n == npoints:
+        return pts
+    rng = np.random.default_rng(seed)
+    return pts[rng.choice(n, size=npoints, replace=n < npoints)]
+
+
+def load_shape(path: str, npoints: int = pa.DEFAULT_NPOINTS, seed: int = 0):
+    """Read a mesh or point cloud; return `npoints` unit-sphere-normalised points + a description."""
+    p = Path(path)
+    ext = p.suffix.lower()
+    src = "point cloud"
+
+    if ext in (".npz", ".npy"):
+        if ext == ".npy":
+            arr = np.load(p)
+        else:
+            d = np.load(p, allow_pickle=True)
+            key = next((k for k in ("points", "xyz", "pos", "vertices") if k in d), None)
+            if key is None:
+                raise gr.Error(f"NPZ must contain points/xyz/pos/vertices — found {list(d.keys())}")
+            arr = d[key]
+        arr = np.asarray(arr, dtype=np.float32)
+        pts = arr.reshape(-1, arr.shape[-1])[:, :3]
+    elif ext in (".txt", ".pts", ".xyz", ".csv", ".asc"):
+        raw = np.loadtxt(p, delimiter="," if ext == ".csv" else None, dtype=np.float32)
+        pts = np.atleast_2d(raw)[:, :3]
+    elif ext in MESH_EXTS:
+        obj = trimesh.load(str(p), process=False)
+        if isinstance(obj, trimesh.Scene):
+            faced = [g for g in obj.geometry.values() if getattr(g, "faces", None) is not None and len(g.faces)]
+            if faced:
+                try:
+                    obj = obj.to_mesh()
+                except Exception:
+                    obj = trimesh.util.concatenate(faced)
+            else:
+                verts = [np.asarray(g.vertices) for g in obj.geometry.values() if hasattr(g, "vertices")]
+                if not verts:
+                    raise gr.Error("No geometry found in this file.")
+                obj = trimesh.PointCloud(np.concatenate(verts, axis=0))
+        if getattr(obj, "faces", None) is not None and len(obj.faces) > 0:
+            np.random.seed(int(seed) % (2**31))
+            pts = np.asarray(trimesh.sample.sample_surface(obj, int(npoints))[0], dtype=np.float32)
+            src = f"mesh, {len(obj.faces):,} faces, surface-sampled"
+        else:
+            pts = np.asarray(obj.vertices, dtype=np.float32)[:, :3]
+    else:
+        raise gr.Error(
+            f"Unsupported file type '{ext}'. Use a mesh (.obj/.glb/.gltf/.stl/.off/.ply) "
+            "or a point cloud (.ply/.npz/.txt/.xyz)."
+        )
+
+    pts = np.ascontiguousarray(pts[np.isfinite(pts).all(axis=1)], dtype=np.float32)
+    if len(pts) < 32:
+        raise gr.Error(f"Only {len(pts)} usable points found — need at least 32.")
+    raw_n = len(pts)
+    pts = _resample(pts, int(npoints), int(seed))
+    return pa.pc_normalize(pts.astype(np.float32)), f"{src}, {raw_n:,} pts → {len(pts):,} used"
+
+
+# --------------------------------------------------------------------------------------
+# Plotting
+# --------------------------------------------------------------------------------------
+
+_AXIS = dict(showbackground=False, showgrid=False, zeroline=False, showticklabels=False, title="")
+
+
+def _style(fig: go.Figure, title: str, height: int) -> go.Figure:
+    fig.update_layout(
+        title=dict(text=title, x=0.02, font=dict(size=12, color="#8a8a8a")),
+        scene=dict(xaxis=_AXIS, yaxis=_AXIS, zaxis=_AXIS, aspectmode="data",
+                   camera=dict(eye=dict(x=1.6, y=1.2, z=1.0))),
+        margin=dict(l=0, r=0, t=28, b=0),
+        height=height,
+        showlegend=len(fig.data) > 1,
+        legend=dict(orientation="h", yanchor="bottom", y=0.0, xanchor="left", x=0.0,
+                    font=dict(color="#8a8a8a", size=11), bgcolor="rgba(0,0,0,0)"),
+        paper_bgcolor="rgba(0,0,0,0)",
+        plot_bgcolor="rgba(0,0,0,0)",
+        font=dict(color="#8a8a8a"),
+    )
+    return fig
+
+
+def plot_raw(points: np.ndarray, title: str, height: int = 300) -> go.Figure:
+    fig = go.Figure(
+        go.Scatter3d(
+            x=points[:, 0], y=points[:, 1], z=points[:, 2], mode="markers",
+            marker=dict(size=1.8, color="#9aa0a6"), name="input", hoverinfo="skip",
+        )
+    )
+    return _style(fig, title, height)
+
+
+def plot_segments(points, pred, names, conf, title: str, height: int = 560) -> go.Figure:
+    fig = go.Figure()
+    for k, name in enumerate(names):
+        m = pred == k
+        if not m.any():
+            continue
+        fig.add_trace(
+            go.Scatter3d(
+                x=points[m, 0], y=points[m, 1], z=points[m, 2], mode="markers",
+                marker=dict(size=2.6, color=PALETTE[k % len(PALETTE)]),
+                name=f"{name} · {int(m.sum())}",
+                customdata=conf[m],
+                hovertemplate=f"<b>{name}</b><br>p=%{{customdata:.2f}}<extra></extra>",
+            )
+        )
+    return _style(fig, title, height)
+
+
+def export_colored_ply(points: np.ndarray, pred: np.ndarray) -> str:
+    rgba = np.zeros((len(points), 4), dtype=np.uint8)
+    rgba[:, 3] = 255
+    for k in range(int(pred.max()) + 1):
+        h = PALETTE[k % len(PALETTE)].lstrip("#")
+        rgba[pred == k, :3] = [int(h[i:i + 2], 16) for i in (0, 2, 4)]
+    f = tempfile.NamedTemporaryFile(suffix="_patchalign3d.ply", delete=False)
+    f.close()
+    trimesh.PointCloud(points, colors=rgba).export(f.name)
+    return f.name
+
+
+# --------------------------------------------------------------------------------------
+# Handlers
+# --------------------------------------------------------------------------------------
+
+
+def preview_shape(shape_file: str, num_points: int = pa.DEFAULT_NPOINTS, seed: int = 0):
+    """Show the uploaded shape as a plain point cloud. CPU only — no GPU needed.
+
+    Args:
+        shape_file: Path to a mesh or point-cloud file.
+        num_points: Number of points to sample for the preview.
+        seed: Sampling seed.
+
+    Returns:
+        An interactive 3D scatter plot of the sampled input points.
+    """
+    if not shape_file:
+        return None
+    points, info = load_shape(shape_file, int(num_points), int(seed))
+    return plot_raw(points, f"Input — {info}")
+
+
+def _parse_labels(labels_text: str):
+    names = [x.strip() for x in (labels_text or "").split(",") if x.strip()]
+    if not names:
+        raise gr.Error("Enter at least one part name, e.g. `back, seat, leg, arm`.")
+    if len(names) > MAX_LABELS:
+        raise gr.Error(f"At most {MAX_LABELS} part queries at a time (got {len(names)}).")
+    return names
+
+
+@spaces.GPU(duration=60)
+def segment(
+    shape_file: str,
+    labels_text: str = "back, seat, leg, arm",
+    num_points: int = pa.DEFAULT_NPOINTS,
+    num_group: int = pa.DEFAULT_NUM_GROUP,
+    group_size: int = pa.DEFAULT_GROUP_SIZE,
+    text_setting: str = "part_only",
+    category: str = "",
+    assign: str = "nearest",
+    tau: float = pa.DEFAULT_TAU,
+    seed: int = 0,
+):
+    """Zero-shot 3D part segmentation of a shape, driven by free-form text part names.
+
+    Args:
+        shape_file: Path to a mesh (.obj/.glb/.gltf/.stl/.off/.ply) or point cloud (.ply/.npz/.txt/.xyz).
+        labels_text: Comma-separated part names to look for, e.g. "back, seat, leg, arm".
+        num_points: Points sampled from the shape (2048 matches the training setting).
+        num_group: Number of patches (furthest-point-sampled centres) the encoder uses.
+        group_size: Points per patch (k-NN neighbourhood size).
+        text_setting: Prompt ensemble — "part_only", "part_plus_cat" or "ensemble".
+        category: Object category used by the "part_plus_cat" / "ensemble" prompts, e.g. "chair".
+        assign: Patch-to-point assignment — "nearest" patch centre, or patch "membership" voting.
+        tau: CLIP temperature used to turn cosine similarities into probabilities.
+        seed: Seed for point / surface sampling.
+
+    Returns:
+        An interactive 3D plot of the segmented shape, the share of points per part,
+        a colour-coded .ply download, and a short run summary.
+    """
+    if not shape_file:
+        raise gr.Error("Upload a 3D shape first, or pick one of the examples below.")
+    names = _parse_labels(labels_text)
+
+    num_points = int(num_points)
+    num_group = max(1, min(int(num_group), num_points))
+    group_size = max(1, min(int(group_size), num_points))
+
+    t0 = time.perf_counter()
+    points, info = load_shape(shape_file, num_points, int(seed))
+    t_load = time.perf_counter() - t0
+
+    t1 = time.perf_counter()
+    pred, probs = pa.segment_point_cloud(
+        points, names, model, proj, text_model, tokenizer, "cuda",
+        category=category or "", text_setting=text_setting, assign=assign,
+        tau=float(tau), num_group=num_group, group_size=group_size,
+    )
+    t_gpu = time.perf_counter() - t1
+
+    conf = probs[np.arange(len(pred)), pred]
+    shares = {name: float((pred == k).mean()) for k, name in enumerate(names)}
+    fig = plot_segments(points, pred, names, conf, "Predicted parts — drag to rotate, scroll to zoom")
+    ply = export_colored_ply(points, pred)
+    summary = (
+        f"**{len(points):,} points → {num_group} patches → {len(names)} text queries**  \n"
+        f"{info} · prompts `{text_setting}`"
+        + (f" · category `{category}`" if category and text_setting != "part_only" else "")
+        + f"  \nload {t_load:.2f}s · inference **{t_gpu:.2f}s** · mean confidence {conf.mean():.2f}"
+    )
+    return fig, shares, ply, summary
+
+
+def apply_preset(preset: str, labels_text: str, category: str):
+    if preset in PRESETS and preset != "— custom —":
+        return PRESETS[preset]
+    return labels_text, category
+
+
+# --------------------------------------------------------------------------------------
+# UI
+# --------------------------------------------------------------------------------------
+
+CSS = """
+#col-container { max-width: 1240px; margin: 0 auto; }
+.dark .gradio-container { color: var(--body-text-color); }
+"""
+
+EXAMPLES = [
+    ["examples/chair.ply", "back, seat, leg"],
+    ["examples/airplane.ply", "body, wing, tail"],
+    ["examples/guitar.ply", "headstock, neck, body"],
+    ["examples/table.ply", "desktop, leg or support, drawer"],
+    ["examples/lamp.ply", "base, lampshade, pole"],
+    ["examples/mug.ply", "handle, cup"],
+    ["examples/bunny.obj", "ear, head, torso, foot"],
+]
+
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="PatchAlign3D") as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+# PatchAlign3D · zero-shot 3D part segmentation
+
+Name the parts you want **in words** and see them highlighted on the 3D shape. One forward pass of a
+point-cloud encoder whose *patch* features are aligned to CLIP text space — no test-time multi-view rendering.
+
+[Paper](https://huggingface.co/papers/2601.02457) · [Code](https://github.com/souhail-hadgi/PatchAlign3D)
+· [Weights](https://huggingface.co/patchalign3d/patchalign3d-encoder)
+· [Project page](https://souhail-hadgi.github.io/patchalign3dsite)
+"""
+        )
+
+        with gr.Row():
+            with gr.Column(scale=2):
+                shape_file = gr.File(
+                    label="3D shape — mesh or point cloud",
+                    file_types=[".obj", ".glb", ".gltf", ".stl", ".off", ".ply",
+                                ".npz", ".npy", ".txt", ".xyz", ".pts"],
+                    type="filepath",
+                )
+                preview = gr.Plot(label="Input")
+                preset = gr.Dropdown(
+                    label="Part-vocabulary preset (fills the box below)",
+                    choices=list(PRESETS.keys()), value="— custom —",
+                )
+                labels_text = gr.Textbox(
+                    label="Part queries (comma-separated)",
+                    value="back, seat, leg, arm",
+                    placeholder="back, seat, leg, arm",
+                    lines=2,
+                )
+                run = gr.Button("Segment", variant="primary")
+
+            with gr.Column(scale=3):
+                plot = gr.Plot(label="Segmentation")
+                summary = gr.Markdown()
+                with gr.Row():
+                    shares = gr.Label(label="Share of points per part", num_top_classes=MAX_LABELS)
+                    ply_out = gr.File(label="Colour-coded point cloud (.ply)")
+
+        with gr.Accordion("Advanced settings", open=False):
+            with gr.Row():
+                num_points = gr.Slider(512, 8192, value=pa.DEFAULT_NPOINTS, step=512, label="Points sampled")
+                num_group = gr.Slider(32, 512, value=pa.DEFAULT_NUM_GROUP, step=32, label="Patches (FPS centres)")
+                group_size = gr.Slider(8, 64, value=pa.DEFAULT_GROUP_SIZE, step=8, label="Points per patch")
+            with gr.Row():
+                text_setting = gr.Radio(
+                    ["part_only", "part_plus_cat", "ensemble"], value="part_only",
+                    label="Prompt ensemble",
+                    info="`part_plus_cat` / `ensemble` also use the object category",
+                )
+                category = gr.Textbox(label="Object category", value="", placeholder="chair")
+            with gr.Row():
+                assign = gr.Radio(["nearest", "membership"], value="nearest", label="Patch → point assignment")
+                tau = gr.Slider(0.01, 1.0, value=pa.DEFAULT_TAU, step=0.01, label="CLIP temperature τ")
+                seed = gr.Number(label="Sampling seed", value=0, precision=0)
+
+        inputs = [shape_file, labels_text, num_points, num_group, group_size,
+                  text_setting, category, assign, tau, seed]
+        outputs = [plot, shares, ply_out, summary]
+
+        gr.Examples(
+            examples=EXAMPLES,
+            inputs=[shape_file, labels_text],
+            outputs=outputs,
+            fn=segment,
+            cache_examples=True,
+            cache_mode="lazy",
+            label="Examples · ShapeNetPart test shapes and the Stanford Bunny mesh",
+        )
+
+        gr.Markdown(
+            """
+### How it works
+Points are centred and scaled to the unit sphere and the Y/Z axes are swapped to match the training
+convention (exactly as in the official `infer.py`). Furthest-point sampling picks patch centres, a k-NN
+neighbourhood around each becomes a patch token, and a 12-layer point transformer produces one feature
+per patch. A learned linear head projects those into the CLIP `ViT-bigG-14 (laion2b_s39b_b160k)` text
+space, where they are matched against the prompt ensemble `{"<part>", "a <part>", "<part> part"}`.
+Each point takes the label of its nearest patch centre.
+
+Every query is *forced* to win somewhere, so asking for a part the shape does not have will still colour
+something — that is expected for open-vocabulary matching. Shapes close to the ShapeNetPart categories
+work best; the Bunny is there to show that arbitrary meshes go through the same path.
+"""
+        )
+
+    shape_file.change(preview_shape, inputs=[shape_file, num_points, seed], outputs=preview,
+                      api_name="preview")
+    preset.change(apply_preset, inputs=[preset, labels_text, category], outputs=[labels_text, category])
+    run.click(segment, inputs=inputs, outputs=outputs, api_name="segment")
+    labels_text.submit(segment, inputs=inputs, outputs=outputs)
+
+if __name__ == "__main__":
+    demo.launch(mcp_server=True)

patchalign3d.py

@@ -0,0 +1,465 @@
+"""Faithful, dependency-light port of the official PatchAlign3D stage-2 inference path.
+
+Source of truth:
+  https://github.com/souhail-hadgi/PatchAlign3D
+    src/models/point_transformer.py   (encoder + patch grouping)
+    src/inference/infer.py            (single-shape inference)
+    src/inference/eval.py             (ShapeNetPart / FAUST evaluation)
+    src/datasets/shapenet.py          (pc_normalize, 2048-point sampling)
+
+Deviations from upstream, all behaviour-preserving:
+  * `pointnet2_ops.furthest_point_sample` -> pure-torch FPS with the same
+    deterministic seeding (start from index 0, squared distances, argmax).
+  * `knn_cuda.KNN(..., transpose_mode=True)` -> pure-torch cdist + topk
+    (ascending distance order, identical semantics).
+  * open_clip `ViT-bigG-14 / laion2b_s39b_b160k` text tower -> the *same*
+    weights served as a HF `CLIPTextModelWithProjection` (verified numerically
+    identical up to fp16 storage rounding), so only the ~1.4 GB text tower is
+    downloaded instead of the full 10 GB two-tower checkpoint.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# --------------------------------------------------------------------------------------
+# Config constants taken verbatim from the reference scripts
+# --------------------------------------------------------------------------------------
+
+TRANS_DIM = 384
+DEPTH = 12
+NUM_HEADS = 6
+ENCODER_DIMS = 256
+DROP_PATH_RATE = 0.1
+
+DEFAULT_NUM_GROUP = 128
+DEFAULT_GROUP_SIZE = 32
+DEFAULT_NPOINTS = 2048
+DEFAULT_TAU = 0.07
+
+CLIP_TEXT_REPO = "stabilityai/stable-diffusion-xl-base-1.0"
+CLIP_TEXT_SUBFOLDER = "text_encoder_2"
+CLIP_TOKENIZER_SUBFOLDER = "tokenizer_2"
+CLIP_TEXT_DIM = 1280  # ViT-bigG-14 joint embedding dim
+
+PART_ONLY_TEMPLATES = ["{}", "a {}", "{} part"]
+PART_PLUS_CAT_TEMPLATES = [
+    "a {} of a {}",
+    "the {} of a {}",
+    "{} of {}",
+    "a {} part of a {}",
+]
+
+
+def clean_text(s: str) -> str:
+    """Upstream `_clean_text`: lowercase, underscores -> spaces, strip punctuation."""
+    s = s.strip().lower().replace("_", " ")
+    out = []
+    for ch in s:
+        out.append(ch if (ch.isalnum() or ch.isspace()) else " ")
+    return " ".join("".join(out).split())
+
+
+# --------------------------------------------------------------------------------------
+# Pure-torch replacements for pointnet2_ops / knn_cuda
+# --------------------------------------------------------------------------------------
+
+
+def furthest_point_sample(xyz: torch.Tensor, npoint: int) -> torch.Tensor:
+    """Iterative FPS matching `pointnet2_ops.furthest_point_sample`.
+
+    Starts from point index 0 and greedily picks the point with the largest
+    squared distance to the already-selected set (exactly what the CUDA kernel
+    does). Returns (B, npoint) long indices.
+    """
+    B, N, _ = xyz.shape
+    device = xyz.device
+    idx = torch.zeros(B, npoint, dtype=torch.long, device=device)
+    dist = torch.full((B, N), 1e10, device=device, dtype=xyz.dtype)
+    farthest = torch.zeros(B, dtype=torch.long, device=device)
+    ar = torch.arange(B, device=device)
+    for i in range(npoint):
+        idx[:, i] = farthest
+        centroid = xyz[ar, farthest, :].view(B, 1, 3)
+        d = ((xyz - centroid) ** 2).sum(-1)
+        dist = torch.minimum(dist, d)
+        farthest = dist.argmax(-1)
+    return idx
+
+
+def fps(data: torch.Tensor, number: int) -> torch.Tensor:
+    """(B, N, 3) -> (B, number, 3) furthest-point-sampled coordinates."""
+    idx = furthest_point_sample(data, number)
+    return torch.gather(data, 1, idx.unsqueeze(-1).expand(-1, -1, data.shape[-1]))
+
+
+def knn_indices(ref: torch.Tensor, query: torch.Tensor, k: int) -> torch.Tensor:
+    """`knn_cuda.KNN(k, transpose_mode=True)(ref, query)[1]`.
+
+    ref: (B, Nr, 3), query: (B, Nq, 3) -> (B, Nq, k) indices into Nr,
+    ordered by ascending distance.
+    """
+    d = torch.cdist(query, ref)  # (B, Nq, Nr)
+    return d.topk(k, dim=-1, largest=False).indices
+
+
+# --------------------------------------------------------------------------------------
+# Point-Transformer encoder (verbatim port of src/models/point_transformer.py)
+# --------------------------------------------------------------------------------------
+
+
+class DropPath(nn.Module):
+    """Stochastic depth. Identity at inference time (which is all we do here)."""
+
+    def __init__(self, drop_prob: float = 0.0):
+        super().__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep = 1.0 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        mask = x.new_empty(shape).bernoulli_(keep).div_(keep)
+        return x * mask
+
+
+class PatchedGroup(nn.Module):
+    """Same as upstream `PatchedGroup`, with FPS/KNN swapped for the torch versions."""
+
+    def __init__(self, num_group: int, group_size: int):
+        super().__init__()
+        self.num_group = num_group
+        self.group_size = group_size
+
+    def forward(self, xyz: torch.Tensor):
+        batch_size, num_points, C = xyz.shape
+        if C > 3:
+            xyz_only = xyz[:, :, :3].contiguous()
+            extra = xyz[:, :, 3:].contiguous()
+        else:
+            xyz_only = xyz.contiguous()
+            extra = None
+
+        center = fps(xyz_only, self.num_group)  # (B, G, 3)
+        idx = knn_indices(xyz_only, center, self.group_size)  # (B, G, M)
+        idx_rel = idx.clone()
+        idx_base = torch.arange(0, batch_size, device=xyz.device).view(-1, 1, 1) * num_points
+        idx_flat = (idx + idx_base).view(-1)
+        neigh_xyz = xyz_only.reshape(batch_size * num_points, -1)[idx_flat, :].view(
+            batch_size, self.num_group, self.group_size, 3
+        )
+        if extra is not None:
+            neigh_extra = extra.reshape(batch_size * num_points, -1)[idx_flat, :].view(
+                batch_size, self.num_group, self.group_size, -1
+            )
+            neighborhood = torch.cat((neigh_xyz - center.unsqueeze(2), neigh_extra), dim=-1)
+        else:
+            neighborhood = neigh_xyz - center.unsqueeze(2)
+        return neighborhood.contiguous(), center.contiguous(), idx_rel
+
+
+class Encoder(nn.Module):
+    def __init__(self, encoder_channel: int, color: bool = False):
+        super().__init__()
+        self.encoder_channel = encoder_channel
+        self.first_conv = nn.Sequential(
+            nn.Conv1d(6 if color else 3, 128, 1),
+            nn.BatchNorm1d(128),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(128, 256, 1),
+        )
+        self.second_conv = nn.Sequential(
+            nn.Conv1d(512, 512, 1),
+            nn.BatchNorm1d(512),
+            nn.ReLU(inplace=True),
+            nn.Conv1d(512, self.encoder_channel, 1),
+        )
+
+    def forward(self, point_groups):
+        bs, g, n, c = point_groups.shape
+        point_groups = point_groups.reshape(bs * g, n, c).permute(0, 2, 1)
+        feature = self.first_conv(point_groups)
+        feature_global = torch.max(feature, 2, keepdim=True)[0]
+        feature_global = feature_global.repeat(1, 1, n)
+        feature = torch.cat([feature_global, feature], 1)
+        feature = self.second_conv(feature)
+        feature = feature.max(dim=2)[0]
+        return feature.reshape(bs, g, self.encoder_channel).contiguous()
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0.0, proj_drop=0.0):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        return self.proj_drop(x)
+
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4.0, qkv_bias=False, qk_scale=None,
+                 drop=0.0, attn_drop=0.0, drop_path=0.0, act_layer=nn.GELU):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                              attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = MLP(in_features=dim, hidden_features=int(dim * mlp_ratio), act_layer=act_layer, drop=drop)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, embed_dim=768, depth=4, num_heads=12, mlp_ratio=4.0, qkv_bias=False,
+                 qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.0):
+        super().__init__()
+
+        def _drop_for_block(i):
+            if isinstance(drop_path_rate, (list, tuple)):
+                return drop_path_rate[i]
+            return drop_path_rate
+
+        self.blocks = nn.ModuleList([
+            Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias,
+                  qk_scale=qk_scale, drop=drop_rate, attn_drop=attn_drop_rate,
+                  drop_path=_drop_for_block(i))
+            for i in range(depth)
+        ])
+
+    def forward(self, x, pos):
+        for blk in self.blocks:
+            x = blk(x + pos)
+        return x
+
+
+class PointTransformer(nn.Module):
+    """Upstream `point_transformer.get_model`."""
+
+    def __init__(self, num_group=DEFAULT_NUM_GROUP, group_size=DEFAULT_GROUP_SIZE, color=False):
+        super().__init__()
+        self.trans_dim = TRANS_DIM
+        self.depth = DEPTH
+        self.num_heads = NUM_HEADS
+        self.encoder_dims = ENCODER_DIMS
+        self.color = color
+        self.group_size = group_size
+        self.num_group = num_group
+
+        self.group_divider = PatchedGroup(num_group=num_group, group_size=group_size)
+        self.encoder = Encoder(encoder_channel=self.encoder_dims, color=color)
+        self.reduce_dim = nn.Linear(self.encoder_dims, self.trans_dim)
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, self.trans_dim))
+        self.cls_pos = nn.Parameter(torch.randn(1, 1, self.trans_dim))
+        self.pos_embed = nn.Sequential(nn.Linear(3, 128), nn.GELU(), nn.Linear(128, self.trans_dim))
+
+        dpr = [x.item() for x in torch.linspace(0, DROP_PATH_RATE, self.depth)]
+        self.blocks = TransformerEncoder(embed_dim=self.trans_dim, depth=self.depth,
+                                        drop_path_rate=dpr, num_heads=self.num_heads)
+        self.norm = nn.LayerNorm(self.trans_dim)
+
+    def set_grouping(self, num_group: int, group_size: int) -> None:
+        self.group_divider.num_group = int(num_group)
+        self.group_divider.group_size = int(group_size)
+
+    def forward_patches(self, pts: torch.Tensor):
+        """pts: (B, C, N) with C >= 3. Returns patch_emb (B, D, G), centers (B, 3, G), idx (B, G, M)."""
+        pts_bn = pts.transpose(-1, -2).contiguous()
+        neighborhood, center, patch_idx = self.group_divider(pts_bn)
+        group_tokens = self.encoder(neighborhood)
+        group_tokens = self.reduce_dim(group_tokens)
+
+        cls_tokens = self.cls_token.expand(group_tokens.size(0), -1, -1)
+        cls_pos = self.cls_pos.expand(group_tokens.size(0), -1, -1)
+        pos = self.pos_embed(center)
+
+        x = torch.cat((cls_tokens, group_tokens), dim=1)
+        pos = torch.cat((cls_pos, pos), dim=1)
+        feature = self.blocks(x, pos)
+        patch_emb = self.norm(feature)[:, 1:, :].transpose(-1, -2).contiguous()
+        patch_centers = center.transpose(-1, -2).contiguous()
+        return patch_emb, patch_centers, patch_idx
+
+
+class PatchToTextProj(nn.Module):
+    def __init__(self, in_dim: int, out_dim: int):
+        super().__init__()
+        self.proj = nn.Linear(in_dim, out_dim)
+
+    def forward(self, patch_emb):
+        x = patch_emb.transpose(1, 2)
+        x = self.proj(x)
+        return F.normalize(x, dim=-1)
+
+
+# --------------------------------------------------------------------------------------
+# Geometry helpers
+# --------------------------------------------------------------------------------------
+
+
+def pc_normalize(pc: np.ndarray) -> np.ndarray:
+    """Upstream `pc_normalize`: centre, then scale to the unit sphere."""
+    centroid = pc.mean(axis=0)
+    pc = pc - centroid
+    m = np.max(np.sqrt((pc ** 2).sum(axis=1)))
+    if m <= 0:
+        m = 1.0
+    return pc / m
+
+
+def prepare_points(points: torch.Tensor) -> torch.Tensor:
+    """Upstream `prepare_points`: (B,N,C) -> (B,C,N) with the Y/Z axes swapped."""
+    if points.ndim != 3:
+        raise ValueError(f"Expected (B,N,C), got {tuple(points.shape)}")
+    pts = points.transpose(2, 1).contiguous()
+    pts[:, [1, 2], :] = pts[:, [2, 1], :]
+    return pts
+
+
+def assign_points_from_patches(points_xyz, patch_centers, patch_logits, patch_idx, mode="nearest"):
+    """Upstream `assign_points_from_patches` (knn_cuda replaced by cdist/argmin)."""
+    B, _, N = points_xyz.shape
+    K = patch_logits.shape[-1]
+    if mode == "membership":
+        point_logits = torch.zeros(B, N, K, device=points_xyz.device, dtype=patch_logits.dtype)
+        counts = torch.zeros(B, N, 1, device=points_xyz.device, dtype=patch_logits.dtype)
+        for b in range(B):
+            idx = patch_idx[b].reshape(-1)
+            src = patch_logits[b].unsqueeze(1).expand_as(patch_idx[b].unsqueeze(-1).expand(-1, -1, K)).reshape(-1, K)
+            point_logits[b].index_add_(0, idx, src)
+            ones = torch.ones(idx.shape[0], 1, device=points_xyz.device, dtype=patch_logits.dtype)
+            counts[b].index_add_(0, idx, ones)
+        return point_logits / counts.clamp_min(1.0)
+    nearest = knn_indices(patch_centers.transpose(1, 2).contiguous(),
+                          points_xyz.transpose(1, 2).contiguous(), 1).squeeze(-1)
+    return patch_logits.gather(1, nearest.unsqueeze(-1).expand(-1, -1, K))
+
+
+# --------------------------------------------------------------------------------------
+# Text side
+# --------------------------------------------------------------------------------------
+
+
+def build_prompts(name: str, category: str, setting: str) -> list[str]:
+    """Prompt ensemble for one part label, mirroring `eval.py:encode_texts`."""
+    nm = clean_text(name)
+    cname = clean_text(category or "")
+    texts: list[str] = []
+    if setting in ("part_plus_cat", "ensemble") and cname:
+        for tpl in PART_PLUS_CAT_TEMPLATES:
+            slots = tpl.count("{}")
+            if slots == 2:
+                texts.append(tpl.format(nm, cname))
+            elif slots == 1:
+                texts.append(tpl.format(f"{cname} {nm}"))
+            else:
+                texts.append(f"{cname} {nm}")
+    if (setting in ("part_only", "ensemble")) or not cname:
+        for tpl in PART_ONLY_TEMPLATES:
+            texts.append(tpl.format(nm) if tpl.count("{}") == 1 else nm)
+    return texts or [nm]
+
+
+@torch.no_grad()
+def encode_labels(names, category, setting, text_model, tokenizer, device) -> torch.Tensor:
+    """One L2-normalised CLIP text embedding per label -> (K, 1280)."""
+    per_label = []
+    for nm in names:
+        prompts = build_prompts(nm, category, setting)
+        toks = tokenizer(prompts, padding="max_length", max_length=tokenizer.model_max_length,
+                         truncation=True, return_tensors="pt").to(device)
+        feat = text_model(**toks).text_embeds.float()
+        feat = F.normalize(feat, dim=-1)
+        per_label.append(F.normalize(feat.mean(dim=0, keepdim=True), dim=-1))
+    return torch.cat(per_label, dim=0)
+
+
+# --------------------------------------------------------------------------------------
+# Checkpoint
+# --------------------------------------------------------------------------------------
+
+
+def load_patchalign3d(ckpt_path: str, num_group=DEFAULT_NUM_GROUP, group_size=DEFAULT_GROUP_SIZE):
+    model = PointTransformer(num_group=num_group, group_size=group_size, color=False)
+    proj = PatchToTextProj(in_dim=TRANS_DIM, out_dim=CLIP_TEXT_DIM)
+    ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+    if "model" in ckpt:
+        res = model.load_state_dict(ckpt["model"], strict=False)
+        print(f"[ckpt] encoder: missing={len(res.missing_keys)} unexpected={len(res.unexpected_keys)}")
+        if res.missing_keys:
+            print("       missing:", res.missing_keys)
+        if res.unexpected_keys:
+            print("       unexpected:", res.unexpected_keys)
+    else:
+        raise RuntimeError("checkpoint has no 'model' entry")
+    if "proj" in ckpt:
+        res = proj.load_state_dict(ckpt["proj"], strict=False)
+        print(f"[ckpt] proj: missing={len(res.missing_keys)} unexpected={len(res.unexpected_keys)}")
+    else:
+        raise RuntimeError("checkpoint has no 'proj' entry")
+    return model.eval(), proj.eval()
+
+
+# --------------------------------------------------------------------------------------
+# End-to-end segmentation
+# --------------------------------------------------------------------------------------
+
+
+@torch.no_grad()
+def segment_point_cloud(points_np, label_names, model, proj, text_model, tokenizer, device,
+                        category="", text_setting="part_only", assign="nearest",
+                        tau=DEFAULT_TAU, num_group=DEFAULT_NUM_GROUP, group_size=DEFAULT_GROUP_SIZE):
+    """points_np: (N,3) float array in original coordinates. Returns (pred, probs)."""
+    model.set_grouping(num_group, group_size)
+    pts = torch.as_tensor(np.ascontiguousarray(points_np[:, :3]), dtype=torch.float32).unsqueeze(0)
+    pts = prepare_points(pts).to(device)
+
+    patch_emb, patch_centers, patch_idx = model.forward_patches(pts)
+    patch_feat = proj(patch_emb)
+    text_feats = encode_labels(label_names, category, text_setting, text_model, tokenizer, device)
+    logits = (patch_feat @ text_feats.t()) / max(float(tau), 1e-6)
+    point_logits = assign_points_from_patches(pts[:, :3, :], patch_centers, logits, patch_idx, mode=assign)
+    probs = point_logits.softmax(dim=-1).squeeze(0)
+    pred = point_logits.argmax(dim=-1).squeeze(0)
+    return pred.cpu().numpy().astype(np.int64), probs.cpu().numpy().astype(np.float32)

requirements.txt

@@ -0,0 +1,8 @@
+torch
+transformers
+numpy
+trimesh
+networkx
+plotly
+safetensors
+pillow