initial test commit

.gitattributes

@@ -0,0 +1,2 @@
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,9 @@
+__pycache__
+lightning_logs*
+logs_*
+.shapenet
+.vscode
+*.ipynb
+*.ini
+scans
+*.venv

Dockerfile

@@ -0,0 +1,34 @@
+FROM pytorch/pytorch:2.5.1-cuda12.1-cudnn9-devel
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PIP_ROOT_USER_ACTION=ignore
+
+RUN apt-get update -qq && \
+    apt-get install -y zip git git-lfs vim libgtk2.0-dev ffmpeg libsm6 libxext6 && \
+    rm -rf /var/cache/apk/*
+
+COPY requirements.txt /workspace
+
+# Activate conda environment and install packages
+RUN conda init bash && \
+    echo "conda activate base" >> ~/.bashrc
+
+SHELL ["conda", "run", "-n", "base", "/bin/bash", "-c"]
+
+RUN pip --no-cache-dir install -r /workspace/requirements.txt
+
+ARG USERNAME=user
+ARG USER_UID=1000
+ARG USER_GID=$USER_UID
+
+# Create the user
+RUN groupadd --gid $USER_GID $USERNAME \
+    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME -s /bin/bash \
+    #
+    # [Optional] Add sudo support. Omit if you don't need to install software after connecting.
+    && apt-get update \
+    && apt-get install -y sudo \
+    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
+    && chmod 0440 /etc/sudoers.d/$USERNAME
+
+WORKDIR /workspaces/pi3detr

LICENSE

@@ -0,0 +1,131 @@
+# PolyForm Noncommercial License 1.0.0
+
+<https://polyformproject.org/licenses/noncommercial/1.0.0>
+
+## Acceptance
+
+In order to get any license under these terms, you must agree
+to them as both strict obligations and conditions to all
+your licenses.
+
+## Copyright License
+
+The licensor grants you a copyright license for the
+software to do everything you might do with the software
+that would otherwise infringe the licensor's copyright
+in it for any permitted purpose.  However, you may
+only distribute the software according to [Distribution
+License](#distribution-license) and make changes or new works
+based on the software according to [Changes and New Works
+License](#changes-and-new-works-license).
+
+## Distribution License
+
+The licensor grants you an additional copyright license
+to distribute copies of the software.  Your license
+to distribute covers distributing the software with
+changes and new works permitted by [Changes and New Works
+License](#changes-and-new-works-license).
+
+## Notices
+
+You must ensure that anyone who gets a copy of any part of
+the software from you also gets a copy of these terms or the
+URL for them above, as well as copies of any plain-text lines
+beginning with `Required Notice:` that the licensor provided
+with the software.  For example:
+
+> Required Notice: Copyright Yoyodyne, Inc. (http://example.com)
+
+## Changes and New Works License
+
+The licensor grants you an additional copyright license to
+make changes and new works based on the software for any
+permitted purpose.
+
+## Patent License
+
+The licensor grants you a patent license for the software that
+covers patent claims the licensor can license, or becomes able
+to license, that you would infringe by using the software.
+
+## Noncommercial Purposes
+
+Any noncommercial purpose is a permitted purpose.
+
+## Personal Uses
+
+Personal use for research, experiment, and testing for
+the benefit of public knowledge, personal study, private
+entertainment, hobby projects, amateur pursuits, or religious
+observance, without any anticipated commercial application,
+is use for a permitted purpose.
+
+## Noncommercial Organizations
+
+Use by any charitable organization, educational institution,
+public research organization, public safety or health
+organization, environmental protection organization,
+or government institution is use for a permitted purpose
+regardless of the source of funding or obligations resulting
+from the funding.
+
+## Fair Use
+
+You may have "fair use" rights for the software under the
+law. These terms do not limit them.
+
+## No Other Rights
+
+These terms do not allow you to sublicense or transfer any of
+your licenses to anyone else, or prevent the licensor from
+granting licenses to anyone else.  These terms do not imply
+any other licenses.
+
+## Patent Defense
+
+If you make any written claim that the software infringes or
+contributes to infringement of any patent, your patent license
+for the software granted under these terms ends immediately. If
+your company makes such a claim, your patent license ends
+immediately for work on behalf of your company.
+
+## Violations
+
+The first time you are notified in writing that you have
+violated any of these terms, or done anything with the software
+not covered by your licenses, your licenses can nonetheless
+continue if you come into full compliance with these terms,
+and take practical steps to correct past violations, within
+32 days of receiving notice.  Otherwise, all your licenses
+end immediately.
+
+## No Liability
+
+***As far as the law allows, the software comes as is, without
+any warranty or condition, and the licensor will not be liable
+to you for any damages arising out of these terms or the use
+or nature of the software, under any kind of legal claim.***
+
+## Definitions
+
+The **licensor** is the individual or entity offering these
+terms, and the **software** is the software the licensor makes
+available under these terms.
+
+**You** refers to the individual or entity agreeing to these
+terms.
+
+**Your company** is any legal entity, sole proprietorship,
+or other kind of organization that you work for, plus all
+organizations that have control over, are under the control of,
+or are under common control with that organization.  **Control**
+means ownership of substantially all the assets of an entity,
+or the power to direct its management and policies by vote,
+contract, or otherwise.  Control can be direct or indirect.
+
+**Your licenses** are all the licenses granted to you for the
+software under these terms.
+
+**Use** means anything you do with the software requiring one
+of your licenses.

app.py

@@ -0,0 +1,1135 @@
+"""
+Gradio + Plotly point cloud viewer for .xyz, .ply and .obj files with PI3DETR model integration.
+
+Features:
+- Upload .xyz (ASCII): one point per line: "x y z" (extra columns are ignored).
+- Upload .ply: Standard PLY format point clouds.
+- Upload .obj: OBJ format with vertices and faces (triangles).
+- Interactive 3D view: orbit, pan, zoom with mouse.
+- Optional: downsample for speed, normalize to unit cube, toggle axes, set point size.
+- Dual view: Input point cloud and model predictions side-by-side.
+- PI3DETR model integration for curve detection.
+- Immediate point cloud rendering on upload.
+"""
+
+import io
+import os
+from typing import List, Dict, Optional
+
+import gradio as gr
+import numpy as np
+import plotly.graph_objects as go
+from plyfile import PlyData
+import pandas
+import torch
+from torch_geometric.data import Data
+import fpsample
+import trimesh  # NEW: for robust mesh loading & surface sampling
+
+# Import PI3DETR modules
+from pi3detr import (
+    build_model,
+    build_model_config,
+    load_args,
+    load_weights,
+)
+from pi3detr.dataset import normalize_and_scale
+
+# Global model cache
+PI3DETR_MODEL = None
+MODEL_STATUS = {"loaded": False, "message": "Model not loaded"}
+
+HOVER_FONT_SIZE = 16  # enlarged hover text size
+FIG_TEMPLATE = "plotly_white"  # global figure template
+PLOT_HEIGHT = 800  # NEW: desired plot height (adjust as needed)
+
+# NEW: demo point cloud file paths (fill these with real .xyz/.ply paths)
+DEMO_POINTCLOUDS = {
+    "Demo 1": "demo_inputs/demo1.xyz",
+    "Demo 2": "demo_inputs/demo2.xyz",
+    "Demo 3": "demo_inputs/demo3.xyz",
+    "Demo 4": "demo_inputs/demo4.xyz",
+    "Demo 5": "demo_inputs/demo5.xyz",
+}
+
+
+def initialize_model(
+    checkpoint_path="checkpoint.ckpt", config_path="configs/pi3detr.yaml"
+):
+    """Initialize the model at startup and store it in the global cache."""
+    global PI3DETR_MODEL, MODEL_STATUS
+    try:
+        args = load_args(config_path) if config_path else {}
+        model_config = build_model_config(args)
+        model = build_model(model_config)
+        load_weights(model, checkpoint_path)
+        model.eval()
+
+        PI3DETR_MODEL = model
+        MODEL_STATUS = {"loaded": True, "message": "Model loaded successfully"}
+        print("PI3DETR model initialized successfully")
+        return True
+    except Exception as e:
+        MODEL_STATUS = {"loaded": False, "message": f"Error loading model: {str(e)}"}
+        print(f"Error initializing PI3DETR model: {e}")
+        return False
+
+
+def read_xyz(file_obj: io.BytesIO) -> np.ndarray:
+    """
+    Parse a .xyz text file from bytes and return Nx3 float32 array.
+    Lines with fewer than 3 numeric values are skipped.
+    Only the first three numeric columns are used.
+    """
+    if file_obj is None:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    # Read bytes → text
+    raw = file_obj.read()
+    try:
+        text = raw.decode("utf-8", errors="ignore")
+    except Exception:
+        text = raw.decode("latin-1", errors="ignore")
+
+    pts = []
+    for line in text.splitlines():
+        line = line.strip()
+        if not line or line.startswith("#"):
+            continue
+        parts = line.replace(",", " ").split()
+        nums = []
+        for p in parts:
+            try:
+                nums.append(float(p))
+            except ValueError:
+                # skip non-numeric tokens
+                pass
+            if len(nums) == 3:
+                break
+        if len(nums) >= 3:
+            pts.append(nums[:3])
+
+    if not pts:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    return np.asarray(pts, dtype=np.float32)
+
+
+def read_ply(file_obj: io.BytesIO) -> np.ndarray:
+    """
+    Parse a .ply file from bytes and return Nx3 float32 array of points.
+    """
+    if file_obj is None:
+        return np.zeros((0, 3), dtype=np.float32)
+
+    try:
+        ply_data = PlyData.read(file_obj)
+        vertex = ply_data["vertex"]
+
+        x = np.asarray(vertex["x"])
+        y = np.asarray(vertex["y"])
+        z = np.asarray(vertex["z"])
+
+        points = np.column_stack([x, y, z]).astype(np.float32)
+        return points
+    except Exception as e:
+        print(f"Error reading PLY file: {e}")
+        return np.zeros((0, 3), dtype=np.float32)
+
+
+def read_obj_and_sample(file_obj: io.BytesIO, display_max_points: int):
+    """Parse OBJ via trimesh and sample up to display_max_points uniformly over the surface."""
+    raw = file_obj.read()
+    # Rewind not strictly needed after read since we don't reuse file_obj
+    try:
+        mesh = trimesh.load(io.BytesIO(raw), file_type="obj", force="mesh")
+    except Exception as e:
+        print(f"trimesh load error: {e}")
+        return (
+            np.zeros((0, 3), dtype=np.float32),
+            np.zeros((0, 3), dtype=np.float32),
+            "OBJ load failure",
+        )
+    # Handle scenes by merging
+    if isinstance(mesh, trimesh.Scene):
+        mesh = trimesh.util.concatenate(tuple(g for g in mesh.geometry.values()))
+    if mesh.is_empty or mesh.vertices.shape[0] == 0:
+        return (
+            np.zeros((0, 3), dtype=np.float32),
+            np.zeros((0, 3), dtype=np.float32),
+            "OBJ: empty mesh",
+        )
+    sample_n = min(display_max_points, max(1, display_max_points))
+    try:
+        sampled = mesh.sample(sample_n)
+    except Exception as e:
+        print(f"Sampling error: {e}")
+        sampled = mesh.vertices
+        if sampled.shape[0] > sample_n:
+            sampled = sampled[:sample_n]
+    sampled = np.asarray(sampled, dtype=np.float32)
+    info = f"OBJ: {mesh.vertices.shape[0]} verts, {len(mesh.faces) if mesh.faces is not None else 0} tris | Surface sampled: {sampled.shape[0]} pts"
+    model_pts = sampled.copy()
+    return model_pts, sampled, info
+
+
+def downsample(pts: np.ndarray, max_points: int) -> np.ndarray:
+    if pts.shape[0] <= max_points:
+        return pts
+    rng = np.random.default_rng(42)
+    idx = rng.choice(pts.shape[0], size=max_points, replace=False)
+    return pts[idx]
+
+
+def make_figure(
+    pts: np.ndarray,
+    point_size: int = 2,
+    show_axes: bool = True,
+    title: str = "",
+    polylines: Optional[List[Dict]] = None,
+) -> go.Figure:
+    """
+    Build a Plotly 3D scatter figure with equal aspect ratio.
+    Optionally includes polylines from model predictions.
+    """
+    if pts.size == 0 and (polylines is None or len(polylines) == 0):
+        fig = go.Figure()
+        fig.update_layout(
+            title="No data to display",
+            template=FIG_TEMPLATE,
+            scene=dict(
+                xaxis_visible=False,
+                yaxis_visible=False,
+                zaxis_visible=False,
+            ),
+            margin=dict(l=0, r=0, t=40, b=0),
+        )
+        return fig
+
+    fig = go.Figure()
+
+    # Add point cloud if available
+    if pts.size > 0:
+        x, y, z = pts[:, 0], pts[:, 1], pts[:, 2]
+        fig.add_trace(
+            go.Scatter3d(
+                x=x,
+                y=y,
+                z=z,
+                mode="markers",
+                marker=dict(
+                    size=max(1, int(point_size)), color="darkgray", opacity=0.2
+                ),
+                hoverinfo="skip",
+                name="Curves",
+                showlegend=False,  # legend hidden
+            )
+        )
+
+    # Define colors for each curve type
+    curve_colors = {
+        "Line": "blue",
+        "Circle": "green",
+        "Arc": "red",
+        "BSpline": "purple",
+    }
+
+    # Add polylines from model predictions if available
+    if polylines:
+        for curve in polylines:
+            points = np.array(curve["points"])
+            if len(points) < 2:
+                continue
+
+            curve_type = curve["type"]
+            curve_id = curve["id"]
+            score = curve["score"]
+
+            # NEW: allow override color if provided (e.g., threshold filtered)
+            color = curve.get("display_color") or curve_colors.get(curve_type, "orange")
+
+            # NEW: support hidden-by-default via legendonly
+            fig.add_trace(
+                go.Scatter3d(
+                    x=points[:, 0],
+                    y=points[:, 1],
+                    z=points[:, 2],
+                    mode="lines",
+                    line=dict(color=color, width=5),
+                    name=f"{curve_type} #{curve_id} ({score:.2f})",
+                    visible=curve.get("visible_state", True),
+                    hoverinfo="text",
+                    text=f"{curve_type} #{curve_id} ({score:.4f})",
+                    showlegend=False,  # hide individual curve legend entries
+                )
+            )
+
+    # Equal aspect ratio using data ranges
+    if pts.size > 0:
+        mins = pts.min(axis=0)
+        maxs = pts.max(axis=0)
+    elif polylines and len(polylines) > 0:
+        # If we only have polylines, calculate range from them
+        all_points = np.vstack([np.array(curve["points"]) for curve in polylines])
+        mins = all_points.min(axis=0)
+        maxs = all_points.max(axis=0)
+    else:
+        mins = np.array([-1, -1, -1])
+        maxs = np.array([1, 1, 1])
+
+    centers = (mins + maxs) / 2.0
+    span = (maxs - mins).max()
+    if span <= 0:
+        span = 1.0
+    half = span / 2.0
+    xrange = [centers[0] - half, centers[0] + half]
+    yrange = [centers[1] - half, centers[1] + half]
+    zrange = [centers[2] - half, centers[2] + half]
+
+    scene_axes = dict(
+        xaxis=dict(range=xrange, visible=show_axes, title="x" if show_axes else ""),
+        yaxis=dict(range=yrange, visible=show_axes, title="y" if show_axes else ""),
+        zaxis=dict(range=zrange, visible=show_axes, title="z" if show_axes else ""),
+        aspectmode="cube",
+    )
+
+    fig.update_layout(
+        title=title,
+        template=FIG_TEMPLATE,
+        showlegend=False,
+        scene=scene_axes,
+        margin=dict(l=0, r=0, t=40, b=0),
+        hoverlabel=dict(font=dict(size=HOVER_FONT_SIZE)),
+        height=PLOT_HEIGHT,  # NEW
+    )
+    return fig
+
+
+def process_model_predictions(data: Data) -> list:
+    """
+    Process model outputs into a format suitable for visualization.
+    """
+    class_names = ["None", "BSpline", "Line", "Circle", "Arc"]
+    polylines = data.polylines.cpu().numpy()
+    curves = []
+
+    # Process detected polylines
+    for i, polyline in enumerate(polylines):
+        cls = data.polyline_class[i].item()
+        score = data.polyline_score[i].item()
+        cls_name = class_names[cls]
+
+        # Skip low-confidence or "None" class predictions
+        if cls == 0:
+            continue
+
+        # Add curve data to results with unique ID
+        curve_data = {
+            "type": cls_name,
+            "id": i + 1,  # 1-based ID for better user experience
+            "index": i,
+            "score": score,
+            "points": polyline,
+        }
+        curves.append(curve_data)
+
+    return curves
+
+
+def process_data_for_model(
+    points: np.ndarray,
+    sample: int = 32768,
+    sample_mode: str = "fps",
+) -> Data:  # CHANGED: removed reduction param
+    """
+    Process and subsample point cloud data using the same approach as predict_pi3detr.py.
+
+    Args:
+        points: Input point cloud as numpy array
+        sample: Number of points to sample
+        sample_mode: Sampling method ("fps", "random", "uniform", "all")
+
+    Returns:
+        Data object ready for model inference
+    """
+    # Convert to torch tensor
+    pos = torch.tensor(points, dtype=torch.float32)
+
+    # Apply sampling strategy
+    if sample_mode == "random":
+        if pos.size(0) > sample:
+            indices = torch.randperm(pos.size(0))[:sample]
+            pos = pos[indices]
+
+    elif sample_mode == "fps":
+        if pos.size(0) > sample:
+            indices = fpsample.bucket_fps_kdline_sampling(pos, sample, h=6)
+            pos = pos[indices]
+
+    elif sample_mode == "uniform":
+        if pos.size(0) > sample:
+            step = max(1, pos.size(0) // sample)
+            pos = pos[::step][:sample]
+
+    elif sample_mode == "all":
+        pass  # Keep all points
+
+    # Create Data object
+    data = Data(pos=pos)
+
+    # Add batch information for single point cloud BEFORE normalization
+    data.batch = torch.zeros(data.pos.size(0), dtype=torch.long)
+    data.batch_size = 1
+
+    # Normalize and scale using PI3DETR's method
+    data = normalize_and_scale(data)
+
+    # Ensure scale and center are proper batch tensors
+    if hasattr(data, "scale") and data.scale.dim() == 0:
+        data.scale = data.scale.unsqueeze(0)
+    if hasattr(data, "center") and data.center.dim() == 1:
+        data.center = data.center.unsqueeze(0)
+
+    return data
+
+
+@torch.no_grad()
+def run_model_inference(
+    model,
+    points: np.ndarray,
+    max_points: int = 32768,
+    sample_mode: str = "fps",
+    num_queries: int = 256,  # NEW
+    snap_and_fit: bool = False,  # NEW
+    iou_filter: bool = False,  # NEW
+) -> list:
+    """Run model inference on the given point cloud (extended with num_queries, snap_and_fit, iou_filter)."""
+    global PI3DETR_MODEL
+    if model is None:
+        model = PI3DETR_MODEL
+    if model is None:
+        return []
+    try:
+        data = process_data_for_model(
+            points, sample=max_points, sample_mode=sample_mode
+        )
+        device = next(model.parameters()).device
+        data = data.to(device)
+
+        if model.num_preds != num_queries:
+            model.set_num_preds(num_queries)
+
+        output = model.predict_step(
+            data,
+            reverse_norm=True,
+            thresholds=None,
+            snap_and_fit=snap_and_fit,  # CHANGED
+            iou_filter=iou_filter,  # CHANGED
+        )
+        result = output[0]
+        curves = process_model_predictions(result)
+        return curves
+    except Exception as e:
+        print(f"Error in model inference: {e}")
+        return []
+
+
+def load_and_process_pointcloud(
+    file: gr.File,
+    max_points: int,
+    point_size: int,
+    show_axes: bool,
+):
+    """
+    Load and process a point cloud from .xyz or .ply file
+    """
+    if file is None:
+        empty_fig = make_figure(np.zeros((0, 3)))
+        return empty_fig, None, None, os.path.basename(file.name) if file else ""
+
+    # Determine file type and read accordingly
+    file_ext = os.path.splitext(file.name)[1].lower()
+
+    # Read file based on extension
+    with open(file.name, "rb") as f:
+        if file_ext == ".xyz":
+            pts = read_xyz(f)
+            mode = "XYZ"
+        elif file_ext == ".ply":
+            pts = read_ply(f)
+            mode = "PLY"
+        elif file_ext == ".obj":
+            model_pts, display_pts, _ = read_obj_and_sample(f, max_points)
+            fig = make_figure(
+                display_pts,
+                point_size=point_size,
+                show_axes=show_axes,
+                title=f"{os.path.basename(file.name)}",
+            )
+            return fig, model_pts, display_pts, os.path.basename(file.name)
+        else:
+            empty_fig = make_figure(np.zeros((0, 3)))
+            return (
+                empty_fig,
+                None,
+                None,
+                "Unsupported file type. Please use .xyz, .ply or .obj.",
+                "",
+            )
+
+    original_n = pts.shape[0]
+
+    # Keep original points for model if normalizing for display
+    model_pts = pts.copy()
+
+    pts = downsample(pts, max_points=max_points)
+    displayed_n = pts.shape[0]
+
+    fig = make_figure(
+        pts,
+        point_size=point_size,
+        show_axes=show_axes,
+        title=f"{os.path.basename(file.name)}",
+    )
+
+    info = f"Loaded ({mode}): {original_n} points"  # | Displayed: {displayed_n} points"
+
+    # RETURN single figure + model/full points + displayed subset
+    return fig, model_pts, pts, os.path.basename(file.name)  # ADDED filename
+
+
+def run_model_prediction(
+    model_pts: np.ndarray,
+    point_size: int,
+    show_axes: bool,
+    model_max_points: int,
+    sample_mode: str,
+    th_bspline: float,
+    th_line: float,
+    th_circle: float,
+    th_arc: float,
+    num_queries: int = 256,
+    snap_and_fit: bool = False,
+    iou_filter: bool = False,
+):  # CHANGED: removed reduction
+    # NOTE: display points now handled outside; keep signature (called before adding display pts state)
+    # (This wrapper kept for backwards compatibility if needed – we adapt below in new unified version)
+    return run_model_prediction_unified(  # type: ignore
+        model_pts,
+        None,
+        point_size,
+        show_axes,
+        model_max_points,
+        sample_mode,
+        th_bspline,
+        th_line,
+        th_circle,
+        th_arc,
+        "",
+        num_queries,
+        snap_and_fit,
+        iou_filter,
+    )
+
+
+def run_model_prediction_unified(
+    model_pts: np.ndarray,
+    display_pts: Optional[np.ndarray],
+    point_size: int,
+    show_axes: bool,
+    model_max_points: int,
+    sample_mode: str,
+    th_bspline: float,
+    th_line: float,
+    th_circle: float,
+    th_arc: float,
+    file_name: str = "",
+    num_queries: int = 256,
+    snap_and_fit: bool = False,
+    iou_filter: bool = False,
+):
+    """
+    Run model inference and apply initial threshold-based coloring.
+    """
+    global PI3DETR_MODEL, MODEL_STATUS
+    if model_pts is None:
+        empty_fig = make_figure(np.zeros((0, 3)))
+        return empty_fig, []
+
+    # Run model inference using cached model
+    curves = []
+    try:
+        if PI3DETR_MODEL is None and not MODEL_STATUS["loaded"]:
+            # Try to initialize model if not already loaded
+            initialize_model()
+
+        if PI3DETR_MODEL is not None:
+            # Run inference with the same settings as predict_pi3detr.py
+            curves = run_model_inference(
+                PI3DETR_MODEL,
+                model_pts,
+                max_points=model_max_points,
+                sample_mode=sample_mode,
+                num_queries=num_queries,  # NEW
+                snap_and_fit=snap_and_fit,  # NEW
+                iou_filter=iou_filter,  # NEW
+            )
+    except Exception:
+        pass
+
+    # NEW: apply thresholds for display (store raw curves separately)
+    thresholds = {
+        "BSpline": th_bspline,
+        "Line": th_line,
+        "Circle": th_circle,
+        "Arc": th_arc,
+    }
+    colored_curves = []
+    for c in curves:
+        c_disp = dict(c)
+        if c["score"] < thresholds.get(c["type"], 0.7):
+            c_disp["visible_state"] = "legendonly"
+        colored_curves.append(c_disp)
+
+    # Use existing displayed subset if provided; else derive lightweight subset
+    if display_pts is None:
+        display_pts = downsample(model_pts, max_points=100000)
+    title = f"{file_name} (curves)" if curves else f"{file_name} (no curves)"
+    fig = make_figure(
+        display_pts,
+        point_size=point_size,
+        show_axes=show_axes,
+        title=title,
+        polylines=colored_curves,
+    )
+    return fig, curves
+
+
+def apply_pointcloud_display_settings(
+    model_pts: np.ndarray,
+    curves: List[Dict],
+    max_points: int,
+    point_size: int,
+    show_axes: bool,
+    th_bspline: float,
+    th_line: float,
+    th_circle: float,
+    th_arc: float,
+    file_name: str,
+):
+    """
+    Apply point cloud display settings without re-running inference.
+    Keeps existing detections and re-applies thresholds.
+    """
+    if model_pts is None:
+        empty_fig = make_figure(np.zeros((0, 3)))
+        return empty_fig, None
+    display_pts = downsample(model_pts, max_points=max_points)
+    if not curves:
+        fig = make_figure(
+            display_pts,
+            point_size=point_size,
+            show_axes=show_axes,
+            title=file_name or "Point Cloud",
+        )
+        return fig, display_pts
+    thresholds = {
+        "BSpline": th_bspline,
+        "Line": th_line,
+        "Circle": th_circle,
+        "Arc": th_arc,
+    }
+    colored_curves = []
+    for c in curves:
+        c_disp = dict(c)
+        if c["score"] < thresholds.get(c["type"], 0.7):
+            c_disp["visible_state"] = "legendonly"
+        colored_curves.append(c_disp)
+    fig = make_figure(
+        display_pts,
+        point_size=point_size,
+        show_axes=show_axes,
+        title=(file_name or "Point Cloud") + " (curves)",
+        polylines=colored_curves,
+    )
+    return fig, display_pts
+
+
+def clear_curves(
+    curves: List[Dict],
+    display_pts: Optional[np.ndarray],
+    model_pts: Optional[np.ndarray],
+    point_size: int,
+    show_axes: bool,
+    file_name: str,
+):
+    """
+    Recolor already inferred curves based on updated thresholds (no re-inference).
+    """
+    if curves is None or model_pts is None or len(curves) == 0:
+        empty_fig = make_figure(
+            display_pts if display_pts is not None else np.zeros((0, 3))
+        )
+        return empty_fig, None
+
+    fig = make_figure(
+        display_pts if display_pts is not None else np.zeros((0, 3)),
+        point_size=point_size,
+        show_axes=show_axes,
+        title=file_name or "Point Cloud",
+        polylines=None,
+    )
+    return fig, None
+
+
+def load_demo_pointcloud(
+    label: str,
+    max_points: int,
+    point_size: int,
+    show_axes: bool,
+):
+    """
+    Load one of the predefined demo point clouds.
+    Clears existing detected curves (curves_state -> None).
+    Also returns a value for the file upload component so the filename shows up.
+    """
+    path = DEMO_POINTCLOUDS.get(label, "")
+    if not path or not os.path.isfile(path):
+        empty_fig = make_figure(np.zeros((0, 3)))
+        return empty_fig, None, None, "", None, None
+    ext = os.path.splitext(path)[1].lower()
+    try:
+        with open(path, "rb") as f:
+            if ext == ".xyz":
+                pts = read_xyz(f)
+            elif ext == ".ply":
+                pts = read_ply(f)
+            elif ext == ".obj":
+                model_pts, display_pts, _ = read_obj_and_sample(
+                    f, min(20000, max_points)
+                )
+                fig = make_figure(
+                    display_pts,
+                    point_size=1,
+                    show_axes=show_axes,
+                    title=f"{os.path.basename(path)} (demo)",
+                )
+                return fig, model_pts, display_pts, os.path.basename(path), None, path
+            else:
+                empty_fig = make_figure(np.zeros((0, 3)))
+                return empty_fig, None, None, "", None, None
+    except Exception:
+        empty_fig = make_figure(np.zeros((0, 3)))
+        return empty_fig, None, None, "", None, None
+    model_pts = pts.copy()
+    pts = downsample(pts, max_points=max_points)
+    fig = make_figure(
+        pts,
+        point_size=1,
+        show_axes=show_axes,
+        title=f"{os.path.basename(path)} (demo)",
+    )
+    return fig, model_pts, pts, os.path.basename(path), None, path
+
+
+# Convenience wrappers for each demo (avoid lambdas for clarity)
+def load_demo1(max_points, point_size, show_axes):
+    return load_demo_pointcloud("Demo 1", max_points, point_size, show_axes)
+
+
+def load_demo2(max_points, point_size, show_axes):
+    return load_demo_pointcloud("Demo 2", max_points, point_size, show_axes)
+
+
+def load_demo3(max_points, point_size, show_axes):
+    return load_demo_pointcloud("Demo 3", max_points, point_size, show_axes)
+
+
+def load_demo4(max_points, point_size, show_axes):  # NEW
+    return load_demo_pointcloud("Demo 4", max_points, point_size, show_axes)
+
+
+def load_demo5(max_points, point_size, show_axes):  # NEW
+    return load_demo_pointcloud("Demo 5", max_points, point_size, show_axes)
+
+
+def build_demo_preview(label: str, max_pts: int = 20000) -> go.Figure:
+    """Create a small preview figure for a demo point cloud (no curves)."""
+    path = DEMO_POINTCLOUDS.get(label, "")
+    if not path or not os.path.isfile(path):
+        return make_figure(np.zeros((0, 3)), title=f"{label}: (missing)")
+    try:
+        ext = os.path.splitext(path)[1].lower()
+        with open(path, "rb") as f:
+            if ext == ".xyz":
+                pts = read_xyz(f)
+            elif ext == ".ply":
+                pts = read_ply(f)
+            elif ext == ".obj":  # UPDATED
+                _, pts, _ = read_obj_and_sample(f, max_pts)
+            else:
+                return make_figure(np.zeros((0, 3)), title=f"{label}: (unsupported)")
+        pts = downsample(pts, max_pts)
+        return make_figure(pts, point_size=1, show_axes=False, title=f"{label} preview")
+    except Exception as e:
+        return make_figure(np.zeros((0, 3)), title=f"{label}: error")
+
+
+def run_model_with_display(
+    model_pts: np.ndarray,
+    max_points: int,
+    point_size: int,
+    show_axes: bool,
+    model_max_points: int,
+    sample_mode: str,
+    th_bspline: float,
+    th_line: float,
+    th_circle: float,
+    th_arc: float,
+    file_name: str = "",
+    num_queries: int = 256,
+    snap_and_fit: bool = False,
+    iou_filter: bool = False,
+):  # CHANGED: removed reduction
+    """
+    Run inference (if model_pts present) then immediately apply current display
+    (max_points/point_size/show_axes) and thresholds. Returns:
+      figure, info_text, curves(list), display_pts
+    """
+    if model_pts is None:
+        empty = make_figure(np.zeros((0, 3)))
+        return empty, None, None
+
+    # Inference first (no display subset passed so it builds from model_pts)
+    fig_infer, curves = run_model_prediction_unified(
+        model_pts,
+        None,
+        point_size,
+        show_axes,
+        model_max_points,
+        sample_mode,
+        th_bspline,
+        th_line,
+        th_circle,
+        th_arc,
+        file_name,
+        num_queries,  # NEW
+        snap_and_fit,  # NEW
+        iou_filter,  # NEW
+    )  # CHANGED: removed reduction
+
+    # Now apply current display settings & thresholds without re-inference
+    fig_final, display_pts = apply_pointcloud_display_settings(
+        model_pts,
+        curves,
+        max_points,
+        point_size,
+        show_axes,
+        th_bspline,
+        th_line,
+        th_circle,
+        th_arc,
+        file_name,
+    )
+    return fig_final, curves, display_pts
+
+
+with gr.Blocks(title="PI3DETR") as demo:
+    gr.Markdown(
+        "# 🥧 PI3DETR: 3D Parametric Curve Inference [CPU-PREVIEW]\n"
+        "An end-to-end deep learning model for **parametric curve inference** in **3D point clouds** and **meshes**.\n"
+        "Upload a `.xyz`, `.ply`, or `.obj` file to explore curve detection."
+    )
+
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(
+                "### 🧩 Supported Inputs\n"
+                "- **Point Clouds:** `.xyz`, `.ply`; **Meshes:** `.obj`\n"
+                "- `Mesh` is surface-sampled using **Max Points (display)** slider."
+            )
+        with gr.Column():
+            gr.Markdown(
+                "### ⚙️ Point Cloud Settings\n"
+                "- Adjust **Max Points**, **point size**, and **axes visibility**.\n"
+                "- Controls visualization of point cloud."
+            )
+        with gr.Column():
+            gr.Markdown(
+                "### 🎯 Confidence Thresholds\n"
+                "- Hover to inspect scores\n."
+                "- Filter curves by **class confidence** interactively"
+            )
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown(
+                "### 🧠 Model Settings\n"
+                "- **Sampling Mode:** Choose downsampling strategy.\n"
+                "- **Model Input Size:** Number of model input points.\n"
+                "- **Queries:** Transformer decoder queries (max. output curves).\n"
+                "- Optional: *Snap&Fit* / *IOU-Filter* post-processing."
+            )
+        with gr.Column():
+            gr.Markdown(
+                "### ⚡ Performance Notes\n"
+                "- Trained on **human-made objects**.\n"
+                "- Optimized for **GPU**; this demo runs on **CPU**.\n"
+                "- For full qualitative performance: \n"
+                "[GitHub → PI3DETR](https://github.com/fafraob/pi3detr)"
+            )
+        with gr.Column():
+            gr.Markdown(
+                "### ▶️ Run Inference\n"
+                "- Click on demo point clouds (from test set) below.\n"
+                "- Press **Run PI3DETR** to execute inference and visualize results."
+            )
+
+    model_pts_state = gr.State(None)
+    display_pts_state = gr.State(None)
+    curves_state = gr.State(None)
+    file_name_state = gr.State("demo_inputs/demo2.xyz")
+    with gr.Row():
+        file_in = gr.File(
+            label="Upload Point Cloud (auto-renders)",
+            file_types=[".xyz", ".ply", ".obj"],
+            type="filepath",
+        )
+    with gr.Row():
+        with gr.Column(scale=1):
+            gr.Markdown("### Point Cloud Settings")
+            max_points = gr.Slider(
+                0,
+                500_000,
+                value=200_000,
+                step=1_000,
+                label="Max points (display)",
+            )
+            point_size = gr.Slider(1, 8, value=1, step=1, label="Point size")
+            show_axes = gr.Checkbox(value=False, label="Show axes")
+
+            gr.Markdown("### Model Settings")
+            sample_mode = gr.Radio(
+                ["fps", "random", "all"],
+                value="fps",
+                label="Main Sampling Method",
+            )
+            model_max_points = gr.Slider(
+                1_000,
+                100_000,
+                value=32768,
+                step=500,
+                label="Downsample to Model Input Size",
+            )
+            num_queries = gr.Slider(  # NEW
+                32,
+                512,
+                value=128,
+                step=1,
+                label="Number of Queries",
+            )
+            with gr.Row():
+                snap_and_fit_chk = gr.Checkbox(value=True, label="Snap&Fit")
+                iou_filter_chk = gr.Checkbox(value=False, label="IOU-Filter")
+
+            # Threshold sliders (no auto-change triggers)
+            gr.Markdown("#### Confidence Thresholds (per class)")
+            th_bspline = gr.Slider(0.0, 1.0, value=0.7, step=0.01, label="BSpline ≥")
+            th_line = gr.Slider(0.0, 1.0, value=0.7, step=0.01, label="Line ≥")
+            th_circle = gr.Slider(0.0, 1.0, value=0.7, step=0.01, label="Circle ≥")
+            th_arc = gr.Slider(0.0, 1.0, value=0.7, step=0.01, label="Arc ≥")
+
+        with gr.Column(scale=1):
+            main_plot = gr.Plot(
+                label="Point Cloud & Curves"
+            )  # height from fig.update_layout(PLOT_HEIGHT)
+
+            run_model_button = gr.Button("Run PI3DETR", variant="primary")
+            clear_curves_button = gr.Button("Clear Curves", variant="secondary")
+
+    # Auto-render point cloud when file is uploaded
+    file_in.change(
+        load_and_process_pointcloud,
+        inputs=[file_in, max_points, point_size, show_axes],
+        outputs=[
+            main_plot,
+            model_pts_state,
+            display_pts_state,
+            file_name_state,
+        ],
+    )
+
+    run_model_button.click(
+        run_model_with_display,
+        inputs=[
+            model_pts_state,
+            max_points,
+            point_size,
+            show_axes,
+            model_max_points,
+            sample_mode,
+            th_bspline,
+            th_line,
+            th_circle,
+            th_arc,
+            file_name_state,
+            num_queries,
+            snap_and_fit_chk,
+            iou_filter_chk,
+        ],
+        outputs=[main_plot, curves_state, display_pts_state],
+    )
+
+    # NEW: auto-apply display & thresholds on interaction (no inference)
+    def _apply_display_wrapper(
+        model_pts,
+        curves,
+        max_points,
+        point_size,
+        show_axes,
+        th_bspline,
+        th_line,
+        th_circle,
+        th_arc,
+        file_name,
+        display_pts_state_value,
+    ):
+        fig, display_pts = apply_pointcloud_display_settings(
+            model_pts,
+            curves,
+            max_points,
+            point_size,
+            show_axes,
+            th_bspline,
+            th_line,
+            th_circle,
+            th_arc,
+            file_name,
+        )
+        return fig, display_pts
+
+    # Point cloud sliders (release) & checkbox (change)
+    for slider in [max_points, point_size]:
+        slider.release(
+            _apply_display_wrapper,
+            inputs=[
+                model_pts_state,
+                curves_state,
+                max_points,
+                point_size,
+                show_axes,
+                th_bspline,
+                th_line,
+                th_circle,
+                th_arc,
+                file_name_state,
+                display_pts_state,
+            ],
+            outputs=[main_plot, display_pts_state],
+        )
+
+    show_axes.change(
+        _apply_display_wrapper,
+        inputs=[
+            model_pts_state,
+            curves_state,
+            max_points,
+            point_size,
+            show_axes,
+            th_bspline,
+            th_line,
+            th_circle,
+            th_arc,
+            file_name_state,
+            display_pts_state,
+        ],
+        outputs=[main_plot, display_pts_state],
+    )
+
+    # Threshold sliders (apply on release)
+    for th in [th_bspline, th_line, th_circle, th_arc]:
+        th.release(
+            _apply_display_wrapper,
+            inputs=[
+                model_pts_state,
+                curves_state,
+                max_points,
+                point_size,
+                show_axes,
+                th_bspline,
+                th_line,
+                th_circle,
+                th_arc,
+                file_name_state,
+                display_pts_state,
+            ],
+            outputs=[main_plot, display_pts_state],
+        )
+
+    clear_curves_button.click(
+        clear_curves,
+        inputs=[
+            curves_state,
+            display_pts_state,
+            model_pts_state,
+            point_size,
+            show_axes,
+            file_name_state,
+        ],
+        outputs=[main_plot, curves_state],
+    )
+
+    # REPLACED demo preview plots + buttons WITH clickable images
+    with gr.Row():
+        gr.Markdown("### Demo Point Clouds (click an image to load)")
+    with gr.Row():
+        # CLEANUP: generate images dynamically for all demos
+        demo_image_components = {}
+        for label in ["Demo 1", "Demo 2", "Demo 3", "Demo 4", "Demo 5"]:  # UPDATED
+            png_path = f"demo_inputs/{label.lower().replace(' ', '')}.png"
+            demo_image_components[label] = gr.Image(
+                value=png_path if os.path.isfile(png_path) else None,
+                label=label,
+                interactive=False,
+            )
+
+    # CLEANUP: map labels to loader functions & attach select handlers
+    _demo_loaders = {
+        "Demo 1": load_demo1,
+        "Demo 2": load_demo2,
+        "Demo 3": load_demo3,
+        "Demo 4": load_demo4,
+        "Demo 5": load_demo5,  # NEW
+    }
+    for label, comp in demo_image_components.items():
+        comp.select(
+            _demo_loaders[label],
+            inputs=[max_points, point_size, show_axes],
+            outputs=[
+                main_plot,
+                model_pts_state,
+                display_pts_state,
+                file_name_state,
+                curves_state,
+                file_in,
+            ],
+        )
+
+    # NEW: auto-load Demo 2 on app start
+    demo.load(
+        load_demo2,
+        inputs=[max_points, point_size, show_axes],
+        outputs=[
+            main_plot,
+            model_pts_state,
+            display_pts_state,
+            file_name_state,
+            curves_state,
+            file_in,
+        ],
+    )
+
+
+if __name__ == "__main__":
+    # Initialize model at startup
+    initialize_model()
+    demo.launch()

configs/pi3detr.yaml

@@ -0,0 +1,77 @@
+### Training parameters
+epochs: 1715
+lr_step: 1250
+lr_warmup_epochs: 15
+lr_warmup_start_factor: 1.0e-6
+lr: 1.0e-4
+batch_size: 8
+batch_size_val: 8
+accumulate_grad_batches: 16
+grad_clip_val: 0.2  # max gradient norm
+to_monitor: "val_seg_iou"
+monitor_mode: "max"
+val_interval: 1
+
+### Model parameters
+model: "pi3detr"
+preencoder_type: "samodule"
+num_features: 0
+weights: ""
+preencoder_lr: 1.0e-4
+freeze_backbone: false
+encoder_dim: 768
+decoder_dim: 768
+num_encoder_layers: 3
+num_decoder_layers: 9
+encoder_dropout: 0.1  # dropout in encoder
+decoder_dropout: 0.1  # dropout in decoder
+num_attn_heads: 8  # number of attention heads
+enc_dim_feedforward: 2048  # dimension of feedforward in encoder
+dec_dim_feedforward: 2048  # dimension of feedforward in decoder
+mlp_dropout: 0.0  # dropout in MLP heads
+num_preds: 128  # num outputs of transformer
+num_classes: 5
+auxiliary_loss: true
+max_points_in_param: 4
+num_transformer_points: 2048  # number of transformer points (needed for some preencoders)
+query_type: "point_fps"
+pos_embed_type: "sine"  # Options: "fourier", "sine"
+class_loss_type: "cross_entropy"
+class_loss_weights: [0.04834912,  0.40329467, 0.09588135, 0.23071379, 0.22176106]
+
+### Curve and validation parameters
+num_curve_points: 64  # must be same as points_per_curve
+num_curve_points_val: 256  # validation curve points
+
+### Loss weights
+loss_weights:
+  loss_class: 1
+  loss_bspline: 1
+  loss_bspline_chamfer: 1
+  loss_line_position: 1
+  loss_line_length: 1
+  loss_line_chamfer: 1
+  loss_circle_position: 1
+  loss_circle_radius: 1
+  loss_circle_chamfer: 1
+  loss_arc: 1
+  loss_arc_chamfer: 1
+  loss_seg: 1
+
+### Cost weights
+cost_weights:
+  cost_class: 1
+  cost_curve: 1
+
+### Dataset parameters
+dataset: "abc_dataset"
+num_workers: 8
+data_root: "/dataset/train"
+data_val_root: "/dataset/val"
+data_test_root: "/dataset/test"
+augment: true
+random_rotate_prob: 1.0
+random_sample_prob: 0.85
+random_sample_bounds: [1.0, 0.2]  # [max, min] fraction of points to keep
+noise_prob: 0.0
+noise_scale: 0.0

configs/pi3detr_k256.yaml

@@ -0,0 +1,77 @@
+### Training parameters
+epochs: 1715
+lr_step: 1250
+lr_warmup_epochs: 15
+lr_warmup_start_factor: 1.0e-6
+lr: 1.0e-4
+batch_size: 8
+batch_size_val: 8
+accumulate_grad_batches: 16
+grad_clip_val: 0.2  # max gradient norm
+to_monitor: "val_seg_iou"
+monitor_mode: "max"
+val_interval: 1
+
+### Model parameters
+model: "pi3detr"
+preencoder_type: "samodule"
+num_features: 0
+weights: ""
+preencoder_lr: 1.0e-4
+freeze_backbone: false
+encoder_dim: 768
+decoder_dim: 768
+num_encoder_layers: 3
+num_decoder_layers: 9
+encoder_dropout: 0.1  # dropout in encoder
+decoder_dropout: 0.1  # dropout in decoder
+num_attn_heads: 8  # number of attention heads
+enc_dim_feedforward: 2048  # dimension of feedforward in encoder
+dec_dim_feedforward: 2048  # dimension of feedforward in decoder
+mlp_dropout: 0.0  # dropout in MLP heads
+num_preds: 256  # num outputs of transformer
+num_classes: 5
+auxiliary_loss: true
+max_points_in_param: 4
+num_transformer_points: 2048  # number of transformer points (needed for some preencoders)
+query_type: "point_fps"
+pos_embed_type: "sine"  # Options: "fourier", "sine"
+class_loss_type: "cross_entropy"
+class_loss_weights: [0.04834912,  0.40329467, 0.09588135, 0.23071379, 0.22176106]
+
+### Curve and validation parameters
+num_curve_points: 64  # must be same as points_per_curve
+num_curve_points_val: 256  # validation curve points
+
+### Loss weights
+loss_weights:
+  loss_class: 1
+  loss_bspline: 1
+  loss_bspline_chamfer: 1
+  loss_line_position: 1
+  loss_line_length: 1
+  loss_line_chamfer: 1
+  loss_circle_position: 1
+  loss_circle_radius: 1
+  loss_circle_chamfer: 1
+  loss_arc: 1
+  loss_arc_chamfer: 1
+  loss_seg: 1
+
+### Cost weights
+cost_weights:
+  cost_class: 1
+  cost_curve: 1
+
+### Dataset parameters
+dataset: "abc_dataset"
+num_workers: 8
+data_root: "/dataset/train"
+data_val_root: "/dataset/val"
+data_test_root: "/dataset/test"
+augment: true
+random_rotate_prob: 1.0
+random_sample_prob: 0.85
+random_sample_bounds: [1.0, 0.2]  # [max, min] fraction of points to keep
+noise_prob: 0.0
+noise_scale: 0.0

pi3detr/__init__.py

@@ -0,0 +1,76 @@
+import argparse
+from typing import Union, Optional
+import torch.nn as nn
+from torch_geometric.data import Dataset
+import inspect
+from .models import ModelConfig
+from .utils import load_args, load_weights
+from .models import PI3DETR
+from .dataset import DatasetConfig, ABCDataset
+
+
+def build_model_config(args: Union[argparse.Namespace, str]) -> ModelConfig:
+    if isinstance(args, str):
+        args = load_args(args)
+
+    # Get required parameters from ModelConfig constructor
+    model_config_signature = inspect.signature(ModelConfig.__init__)
+    required_params = [
+        param for param in model_config_signature.parameters.keys() if param != "self"
+    ]
+
+    for param in required_params:
+        if not hasattr(args, param):
+            print(f"ERROR: Parameter '{param}' has to be specified in the arguments")
+            raise ValueError(f"Missing required parameter: {param}")
+
+    # Create model config with all parameters from args
+    model_config_args = {param: getattr(args, param) for param in required_params}
+    model_config = ModelConfig(**model_config_args)
+
+    print(model_config)
+    return model_config
+
+
+def build_dataset_config(
+    args: Union[argparse.Namespace, str], data_root: str, augment: bool
+) -> DatasetConfig:
+    if isinstance(args, str):
+        args = load_args(args)
+
+    # Get required parameters from DatasetConfig constructor (excluding root and augment)
+    dataset_config_signature = inspect.signature(DatasetConfig.__init__)
+    required_params = [
+        param
+        for param in dataset_config_signature.parameters.keys()
+        if param not in ["self", "root", "augment"]
+    ]
+
+    for param in required_params:
+        if not hasattr(args, param):
+            print(f"ERROR: Parameter '{param}' has to be specified in the arguments")
+            raise ValueError(f"Missing required parameter: {param}")
+
+    # Create dataset config with parameters from args plus root and augment
+    dataset_config_args = {param: getattr(args, param) for param in required_params}
+    dataset_config = DatasetConfig(
+        root=data_root, augment=augment, **dataset_config_args
+    )
+
+    print(dataset_config)
+    return dataset_config
+
+
+def build_dataset(config: DatasetConfig) -> Dataset:
+    if config.dataset == "abc_dataset":
+        return ABCDataset(config)
+    else:
+        raise ValueError(f"Unknown dataset {config.dataset}")
+
+
+def build_model(config: ModelConfig) -> nn.Module:
+    print(f"Model: {config.model}")
+    if config.model == "pi3detr":
+        return PI3DETR(config)
+    else:
+        raise ValueError(f"Unknown model {config.model}")

pi3detr/dataset/__init__.py

@@ -0,0 +1,7 @@
+from .abc_dataset import ABCDataset, DatasetConfig
+from .point_cloud_transforms import (
+    normalize_and_scale,
+    normalize_and_scale_with_params,
+    reverse_normalize_and_scale,
+    reverse_normalize_and_scale_with_params,
+)

pi3detr/dataset/abc_dataset.py

@@ -0,0 +1,159 @@
+import random
+import torch
+import numpy as np
+from typing import Union
+from pathlib import Path
+from torch_geometric.data import Dataset
+from torch_geometric.data.data import BaseData
+from dataclasses import dataclass
+from typing import Callable
+import torch_geometric.transforms as T
+
+from .point_cloud_transforms import (
+    random_rotate,
+    normalize_and_scale,
+    add_noise,
+    subsample,
+)
+
+
+@dataclass
+class DatasetConfig:
+    dataset: str
+    root: str
+    augment: bool = False
+    random_rotate_prob: float = 1
+    random_sample_prob: float = 0.5
+    random_sample_bounds: tuple[float, float] = (1, 0.5)
+    noise_prob: float = 0
+    noise_scale: float = 0
+
+
+class ABCDataset(Dataset):
+    def __init__(
+        self,
+        config: DatasetConfig,
+    ) -> None:
+        self.file_names = self._read_file_names(config.root)
+        self.config = config
+        super().__init__(
+            config.root,
+            None,
+            None,
+            None,
+        )
+
+    @property
+    def raw_dir(self) -> str:
+        return self.root
+
+    @property
+    def raw_file_names(self) -> Union[str, list[str], tuple]:
+        return self.file_names
+
+    @property
+    def processed_file_names(self) -> Union[str, list[str], tuple]:
+        return [f"{file_name}.pt" for file_name in self.file_names]
+
+    def process(self) -> None:
+        print("Should already be processed.")
+
+    def get(self, idx: int) -> BaseData:
+
+        data = torch.load(
+            Path(self.processed_dir) / f"{self.raw_file_names[idx]}.pt",
+            weights_only=False,
+        )
+        data["pos"] = data["pos"].to(torch.float32)
+
+        augment = self.config.augment
+        if augment and random.random() < self.config.noise_prob:
+            sigma = (
+                np.max(
+                    np.max(data.pos.cpu().numpy(), axis=0)
+                    - np.min(data.pos.cpu().numpy(), axis=0)
+                )
+                / self.config.noise_scale
+            )
+            noise = torch.tensor(
+                np.random.normal(loc=0, scale=sigma, size=data.pos.shape),
+                dtype=data.pos.dtype,
+                device=data.pos.device,
+            )
+            data.pos += noise
+
+        if not hasattr(data, "real_scale") or not hasattr(data, "real_center"):
+            data.real_center = torch.zeros(3)
+            data.real_scale = torch.tensor(1.0)
+
+        if augment and random.random() < self.config.random_sample_prob:
+            data = subsample(
+                data,
+                *self.config.random_sample_bounds,
+                max_points=None,
+                extra_fields=["y_seg", "y_seg_cls"],
+            )
+
+        extra_fields = [
+            "y_curve_64",
+            "bspline_params",
+            "line_params",
+            "circle_params",
+            "arc_params",
+        ]
+        if augment and random.random() < self.config.random_rotate_prob:
+            data = random_rotate(data, 180, axis=0, extra_fields=extra_fields)
+            data = random_rotate(data, 180, axis=1, extra_fields=extra_fields)
+            data = random_rotate(data, 180, axis=2, extra_fields=extra_fields)
+
+        line_direction = data.line_params[:, 1]
+        circle_normal = data.circle_params[:, 1]
+
+        data = normalize_and_scale(
+            data,
+            extra_fields=extra_fields,
+        )
+        # normal vecotrs shouldn't change
+        data.line_params[:, 1] = line_direction
+        data.circle_params[:, 1] = circle_normal
+        # manually adjust length and radius
+        data.line_length = data.line_length * data.scale
+        data.circle_radius = data.circle_radius * data.scale
+
+        data.y_params = torch.zeros(data.num_curves, 12, dtype=torch.float32)
+        for i in range(data.num_curves):
+            if data.y_cls[i] == 1:
+                # B-spline
+                # P0, P1, P2, P3
+                data.y_params[i][:12] = data.bspline_params[i].reshape(-1)
+            elif data.y_cls[i] == 2:
+                # Line
+                # midpoint, normal, length
+                data.y_params[i][:3] = data.line_params[i][0].reshape(-1)
+                data.y_params[i][3:6] = line_direction[i].reshape(-1)
+                data.y_params[i][6] = data.line_length[i]  # already adjusted above
+            elif data.y_cls[i] == 3:
+                # Circle
+                # center, normal, radius
+                data.y_params[i][:3] = data.circle_params[i][0].reshape(-1)
+                data.y_params[i][3:6] = circle_normal[i].reshape(-1)
+                data.y_params[i][6] = data.circle_radius[i]  # already adjusted above
+            elif data.y_cls[i] == 4:
+                # Arc
+                # midpoint, start, end
+                data.y_params[i][:9] = data.arc_params[i].reshape(-1)
+        data.filename = self.raw_file_names[idx]
+
+        return data
+
+    def len(self) -> int:
+        return len(self.processed_file_names)
+
+    def _read_file_names(self, root: str) -> list[Path]:
+        return sorted(
+            [
+                fp.stem
+                for fp in Path(root).joinpath("processed").glob(f"*.pt")
+                if "pre_" not in fp.stem
+            ]
+        )

pi3detr/dataset/point_cloud_transforms.py

@@ -0,0 +1,259 @@
+import math
+import random
+import torch
+import numpy as np
+from typing import Optional
+import torch_geometric.transforms as T
+from torch_geometric.data import Data
+
+
+def subsample(
+    data: Data,
+    upper_bound: float = 1.0,
+    lower_bound: float = 0.5,
+    max_points: Optional[int] = None,
+    extra_fields: list[str] = [],
+) -> Data:
+    r"""Subsamples the point cloud to a random number of points within the
+    range :obj:`[lower_bound, upper_bound]` (functional name: :obj:`subsample`).
+    """
+    if data.pos.size(0) == 0:
+        return data
+    num_points = int(random.uniform(lower_bound, upper_bound) * data.pos.size(0))
+    if max_points is not None:
+        num_points = min(num_points, max_points)
+    idx = torch.randperm(data.pos.size(0))[:num_points]
+    data.pos = data.pos[idx]
+    for field in extra_fields:
+        if hasattr(data, field):
+            setattr(data, field, getattr(data, field)[idx])
+    return data
+
+
+def numpy_normalize_and_scale(xyz: np.ndarray) -> tuple[np.ndarray, float, float]:
+    r"""Normalizes the point cloud in such a way that the points are centered
+    around the origin and are within the interval :math:`[-1, 1]` (functional
+    name: :obj:`normalize`).
+    """
+    center = xyz.mean(0)
+    scale = (1 / np.max(np.abs(xyz - center))) * 0.999999
+    xyz = numpy_normalize_and_scale_with_params(xyz, center, scale)
+    return xyz, center, scale
+
+
+def numpy_normalize_and_scale_with_params(
+    xyz: np.ndarray, center: np.ndarray, scale: float
+) -> np.ndarray:
+    r"""Normalizes the point cloud in such a way that the points are centered
+    around the origin and are within the interval :math:`[-1, 1]` (functional
+    name: :obj:`normalize`).
+    """
+    if xyz.size == 0:
+        return xyz
+    shape = xyz.shape
+    return ((xyz.reshape(-1, shape[-1]) - center) * scale).reshape(shape)
+
+
+def normalize_and_scale(data: Data, extra_fields: list[str] = []) -> Data:
+    r"""Centers and normalizes the given fields to the interval :math:`[-1, 1]`
+    (functional name: :obj:`normalize_scale`).
+    """
+    if data.pos.size(0) == 0:
+        data.center = torch.empty(0)
+        data.scale = torch.empty(0)
+        return data
+    # center the pos points
+    center = data.pos.mean(dim=-2, keepdim=True)
+    # scale the pos points
+    scale = (1 / (data.pos - center).abs().max()) * 0.999999
+
+    return normalize_and_scale_with_params(data, center, scale, extra_fields)
+
+
+def reverse_normalize_and_scale(data: Data, extra_fields: list[str] = []) -> Data:
+    r"""Reverses the centering and normalization of the given fields
+    (functional name: :obj:`reverse_normalize_scale`).
+    """
+    assert hasattr(data, "center") and hasattr(
+        data, "scale"
+    ), "Data object does not contain the center and scale attributes."
+    return reverse_normalize_and_scale_with_params(
+        data, data.center, data.scale, extra_fields
+    )
+
+
+def normalize_and_scale_with_params(
+    data: Data, center: torch.Tensor, scale: torch.Tensor, extra_fields: list[str] = []
+) -> Data:
+    if data.pos.size(0) == 0:
+        data.center = torch.empty(0)
+        data.scale = torch.empty(0)
+        return data
+    data.pos = (data.pos - center) * scale
+    for field in extra_fields:
+        if hasattr(data, field):
+            shape = getattr(data, field).size()
+            setattr(
+                data,
+                field,
+                (getattr(data, field).reshape(-1, shape[-1]) - center) * scale,
+            )
+            setattr(data, field, getattr(data, field).reshape(shape))
+    data.center = center
+    data.scale = scale
+    return data
+
+
+def reverse_normalize_and_scale_with_params(
+    data: Data, center: torch.Tensor, scale: torch.Tensor, extra_fields: list[str] = []
+) -> Data:
+    r"""Reverses the centering and normalization of the given fields
+    (functional name: :obj:`reverse_normalize_scale`).
+    """
+    # Reverse the scaling and centering of the pos points
+    data.pos = data.pos / scale + center
+
+    for field in extra_fields:
+        if hasattr(data, field):
+            shape = getattr(data, field).size()
+            setattr(
+                data,
+                field,
+                (getattr(data, field).reshape(-1, shape[-1]) / scale) + center,
+            )
+            setattr(data, field, getattr(data, field).reshape(shape))
+    data.center = torch.empty(0)
+    data.scale = torch.empty(0)
+    return data
+
+
+def reverse_normalize_and_scale_with_params(
+    data: Data, center: torch.Tensor, scale: torch.Tensor, extra_fields: list[str] = []
+) -> Data:
+    r"""Reverses the centering and normalization of the given fields
+    (functional name: :obj:`reverse_normalize_scale`).
+    """
+    # Reverse the scaling and centering of the pos points
+    data.pos = data.pos / scale + center
+
+    for field in extra_fields:
+        if hasattr(data, field):
+            shape = getattr(data, field).size()
+            setattr(
+                data,
+                field,
+                (getattr(data, field).reshape(-1, shape[-1]) / scale) + center,
+            )
+            setattr(data, field, getattr(data, field).reshape(shape))
+    data.center = torch.empty(0)
+    data.scale = torch.empty(0)
+    return data
+
+
+def random_rotate(
+    data: Data, degrees: float, axis: int, extra_fields: list[str] = []
+) -> Data:
+    r"""Rotates the object around the origin by a random angle within the
+    range :obj:`[-degrees, degrees]` (functional name: :obj:`random_rotate
+    `).
+    """
+    if data.pos.size(0) == 0:
+        return data
+    return rotate_with_params(
+        data, random.uniform(-degrees, degrees), axis, extra_fields
+    )
+
+
+def rotate_with_params(
+    data: Data, degrees: float, axis: int = 0, extra_fields: list[str] = []
+) -> Data:
+    r"""Rotates the object around the origin by a given angle
+    (functional name: :obj:`rotate`).
+    """
+    angle = math.pi * degrees / 180.0
+    if data.pos.size(0) == 0:
+        return data
+    sin, cos = math.sin(angle), math.cos(angle)
+    if data.pos.size(-1) == 2:
+        matrix = torch.tensor([[cos, sin], [-sin, cos]])
+    else:
+        if axis == 0:
+            matrix = torch.tensor([[1, 0, 0], [0, cos, sin], [0, -sin, cos]])
+        elif axis == 1:
+            matrix = torch.tensor([[cos, 0, -sin], [0, 1, 0], [sin, 0, cos]])
+        else:
+            matrix = torch.tensor([[cos, sin, 0], [-sin, cos, 0], [0, 0, 1]])
+
+    matrix_dtype = matrix.to(data.pos.dtype)
+    matrix = matrix.to(matrix_dtype)
+
+    data.pos = data.pos @ matrix.t()
+    for field in extra_fields:
+        if hasattr(data, field):
+            shape = getattr(data, field).size()
+            # get dtype of field
+            dtype = getattr(data, field).dtype
+
+            matrix_dtype = matrix.to(dtype)
+            setattr(data, field, getattr(data, field) @ matrix_dtype.t())
+            setattr(data, field, getattr(data, field).reshape(shape))
+    setattr(data, f"rotated_{axis}", degrees)
+    return data
+
+
+def reverse_rotate(data: Data, axis: int = 0, extra_fields: list[str] = []) -> Data:
+    r"""Reverses the rotation of the object around the origin
+    (functional name: :obj:`reverse_rotate`).
+    """
+    if not hasattr(data, f"rotated_{axis}"):
+        return data
+    return rotate_with_params(
+        data, -getattr(data, f"rotated_{axis}"), axis, extra_fields
+    )
+
+
+def add_noise(data: Data, std: float) -> Data:
+    r"""Adds Gaussian noise to the node features (functional name:
+    :obj:`add_noise`).
+    """
+    if data.pos.size(0) == 0:
+        return data
+    noise = torch.randn_like(data.pos) * std
+    data.pos = data.pos + noise
+    data.noise = noise
+    return data
+
+
+def remove_noise(data: Data) -> Data:
+    r"""Removes the noise from the node features (functional name:
+    :obj:`remove_noise`).
+    """
+    assert hasattr(data, "noise"), "Data object does not contain the noise attribute."
+    data.pos = data.pos - data.noise
+    del data.noise
+    return data
+
+
+def custom_normalize_and_scale(
+    data: Data, p1: torch.Tensor, p2: torch.Tensor, extra_fields: list[str] = []
+) -> Data:
+    r"""Normalizes the point cloud in such a way that after the transformation
+    p1 is at (0,0,0) and p2 at (1,1,1)` (functional name:
+    :obj:`normalize`).
+    """
+    assert p1.size() == p2.size() == (3,), "Invalid interval."
+    if data.pos.size(0) == 0:
+        return data
+    data.pos = (data.pos - p1) / (p2 - p1)
+    for field in extra_fields:
+        if hasattr(data, field):
+            shape = getattr(data, field).size()
+            setattr(
+                data,
+                field,
+                (getattr(data, field).reshape(-1, shape[-1]) - p1) / (p2 - p1),
+            )
+            setattr(data, field, getattr(data, field).reshape(shape))
+    data.p1 = p1
+    data.p2 = p2
+    return data

pi3detr/dataset/utils.py

@@ -0,0 +1,75 @@
+import numpy as np
+import json
+import open3d as o3d
+
+
+def read_xyz_file(file_path: str, column_idxs: list[int] = [0, 1, 2]) -> np.ndarray:
+    """Reads a point cloud from a .xyz file.
+
+    Args:
+        file_path (str): Path to the .xyz file.
+        column_idxs (list[int], optional): Indices of the columns to read. Defaults to [0,1,2].
+
+    Returns:
+        np.ndarray: Point cloud as a numpy array.
+    """
+    return np.loadtxt(file_path, usecols=column_idxs)
+
+
+def read_curve_file(file_path: str) -> tuple[np.ndarray]:
+    with open(file_path, "r") as f:
+        data = json.load(f)
+    return np.array(data["linear"]), np.array(data["bezier"])
+
+
+def read_polyline_file(file_path: str, sep: str = ",") -> np.ndarray:
+    polylines = []
+    with open(file_path, "r") as f:
+        polyline = []
+        for line in f:
+            if line == "\n":
+                polylines.append(polyline)
+                polyline = []
+            else:
+                point = [float(x) for x in line.split(sep)]
+                polyline.append(point)
+        if polyline:
+            polylines.append(polyline)
+        return np.array(polylines)
+
+
+def voxel_down_sample(xyz: np.ndarray, voxel_size: float) -> np.ndarray:
+    """Downsamples a point cloud using voxel grid downsampling.
+
+    Args:
+        xyz (np.ndarray): Point cloud.
+        voxel_size (float): Voxel size.
+
+    Returns:
+        np.ndarray: Downsampled point cloud.
+    """
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(xyz)
+    downpcd = pcd.voxel_down_sample(voxel_size=voxel_size)
+    return np.asarray(downpcd.points)
+
+
+def filter_normals(
+    xyz: np.ndarray, radius: float, max_nn: float, threshold: float
+) -> np.ndarray:
+    """Filters normals of a point cloud.
+
+    Args:
+        xyz (np.ndarray): Point cloud.
+        threshold (float, optional): Threshold for filtering normals.
+
+    Returns:
+        np.ndarray: Filtered point cloud.
+    """
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(xyz)
+    pcd.estimate_normals(
+        search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=max_nn)
+    )
+    new_pts = np.asarray(pcd.points)[np.abs(np.asarray(pcd.normals)[:, 2]) < threshold]
+    return new_pts

pi3detr/evaluation/__init__.py

@@ -0,0 +1 @@
+from .abc_metrics import ChamferMAP, ChamferIntervalMetric

pi3detr/evaluation/abc_metrics.py

@@ -0,0 +1,349 @@
+import torch
+import numpy as np
+from torchmetrics import Metric
+from torchmetrics.functional import average_precision
+from scipy.spatial import KDTree
+
+
+def calc_chamfer_distance(pred_points, gt_points):
+    """Calculate chamfer distance and bidirectional hausdorff distance."""
+    if len(pred_points) == 0 or len(gt_points) == 0:
+        return float("inf"), float("inf")
+
+    tree_pred = KDTree(pred_points)
+    tree_gt = KDTree(gt_points)
+
+    dist_pred2gt, _ = tree_gt.query(pred_points)
+    dist_gt2pred, _ = tree_pred.query(gt_points)
+
+    chamfer_dist = np.mean(dist_pred2gt**2) + np.mean(dist_gt2pred**2)
+    bhaussdorf_dist = (dist_pred2gt.max() + dist_gt2pred.max()) / 2
+
+    return chamfer_dist, bhaussdorf_dist
+
+
+class ChamferMAP(Metric):
+    def __init__(self, chamfer_thresh=0.05, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.chamfer_thresh = chamfer_thresh
+        self.class_names = {
+            1: "mAP_bspline",
+            2: "mAP_line",
+            3: "mAP_circle",
+            4: "mAP_arc",
+        }
+
+        self.add_state("all_scores", default=[], dist_reduce_fx="cat")
+        self.add_state("all_matches", default=[], dist_reduce_fx="cat")
+        self.add_state("all_classes", default=[], dist_reduce_fx="cat")
+
+    def pairwise_chamfer_distance_batch(self, pred, gt):
+        """
+        pred: [P, 64, 3]
+        gt:   [G, 64, 3]
+        returns: [P, G] chamfer distances
+        """
+        P, G = pred.size(0), gt.size(0)
+
+        # Reshape for pairwise comparison
+        pred_exp = pred.unsqueeze(1)  # [P, 1, 64, 3]
+        gt_exp = gt.unsqueeze(0)  # [1, G, 64, 3]
+
+        # Compute pairwise distances between points
+        dists = torch.cdist(pred_exp, gt_exp, p=2)  # [P, G, 64, 64]
+
+        a2b = dists.min(dim=3).values.mean(dim=2)  # [P, G]
+        b2a = dists.min(dim=2).values.mean(dim=2)  # [P, G]
+
+        return a2b + b2a  # [P, G]
+
+    def update(self, outputs, batch):
+        B = outputs["pred_class"].shape[0]
+        y_curves = batch.y_curve_64  # [total_gt, 64, 3]
+        y_cls = batch.y_cls  # [total_gt]
+        num_curves_per_batch = batch.num_curves.tolist()
+
+        gt_splits = torch.split(y_curves, num_curves_per_batch, dim=0)
+        cls_splits = torch.split(y_cls, num_curves_per_batch, dim=0)
+
+        pred_classes_all = outputs["pred_class"].softmax(dim=-1)  # [B, N, C]
+        for b in range(B):
+            pred_classes = pred_classes_all[b]  # [N, C]
+
+            preds_all = {
+                1: outputs["pred_bspline_points"][b],  # [N, 64, 3]
+                2: outputs["pred_line_points"][b],
+                3: outputs["pred_circle_points"][b],
+                4: outputs["pred_arc_points"][b],
+            }
+
+            for cls in self.class_names.keys():
+                pred_points = preds_all[cls]  # [P, 64, 3]
+                scores = pred_classes[:, cls]  # [P]
+
+                gt_points = gt_splits[b][cls_splits[b] == cls]  # [G, 64, 3]
+                if gt_points.size(0) == 0:
+                    self.all_scores.append(scores)
+                    self.all_matches.append(torch.zeros_like(scores))
+                    self.all_classes.append(torch.full_like(scores, cls))
+                    continue
+
+                chamfer = self.pairwise_chamfer_distance_batch(pred_points, gt_points)
+                used_gt = torch.zeros(
+                    gt_points.size(0), dtype=torch.bool, device=pred_points.device
+                )
+                matches = torch.zeros(pred_points.size(0), device=pred_points.device)
+
+                for i in range(pred_points.size(0)):
+                    dists = chamfer[i]
+                    min_dist, min_idx = dists.min(0)
+                    if min_dist < self.chamfer_thresh and not used_gt[min_idx]:
+                        matches[i] = 1.0
+                        used_gt[min_idx] = True
+
+                self.all_scores.append(scores)
+                self.all_matches.append(matches)
+                self.all_classes.append(torch.full_like(matches, cls))
+
+    def compute(self):
+        if not self.all_scores:
+            return {cls_name: 0.0 for cls_name in self.class_names.values()}
+
+        scores = torch.cat(self.all_scores)
+        matches = torch.cat(self.all_matches)
+        classes = torch.cat(self.all_classes)
+
+        result = {}
+        ap_values = []
+
+        for cls in self.class_names.keys():
+            mask = classes == cls
+            if mask.sum() == 0 or torch.sum(matches[mask]) == 0:
+                ap = torch.tensor(0.0, device=self.device)
+            else:
+                ap = average_precision(
+                    scores[mask], matches[mask].to(torch.int32), task="binary"
+                )
+            result[self.class_names[cls]] = ap.item()
+            ap_values.append(ap)
+
+        result["mAP"] = torch.stack(ap_values).mean().item()
+        return result
+
+
+class ChamferIntervalMetric(Metric):
+
+    def __init__(self, interval=0.01, map_cd_thresh=0.005, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.interval = interval
+        self.add_state("total_cd", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_cd_sq", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_bhd", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total_bhd_sq", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("count", default=torch.tensor(0), dist_reduce_fx="sum")
+        self.add_state("valid_count", default=torch.tensor(0), dist_reduce_fx="sum")
+
+        self.map_cd_thresh = map_cd_thresh
+        self.map_cls_names = {
+            1: "mAP_bspline",
+            2: "mAP_line",
+            3: "mAP_circle",
+            4: "mAP_arc",
+        }
+
+        self.add_state("map_all_scores", default=[], dist_reduce_fx="cat")
+        self.add_state("map_all_matches", default=[], dist_reduce_fx="cat")
+        self.add_state("map_all_classes", default=[], dist_reduce_fx="cat")
+
+    def sample_curve_by_interval(self, points, interval, force_last=False):
+        """Sample points along the curve at fixed length `interval`."""
+        if len(points) < 2:
+            return points
+
+        edges = np.array([[j, j + 1] for j in range(len(points) - 1)])
+        edge_lengths = np.linalg.norm(points[edges[:, 1]] - points[edges[:, 0]], axis=1)
+
+        samples = [points[0]]
+        distance_accum = 0.0
+        next_sample_dist = interval
+        edge_index = 0
+
+        while edge_index < len(edges):
+            p0 = points[edges[edge_index, 0]]
+            p1 = points[edges[edge_index, 1]]
+            edge_vec = p1 - p0
+            edge_len = np.linalg.norm(edge_vec)
+            if edge_len == 0:
+                edge_index += 1
+                continue
+
+            while distance_accum + edge_len >= next_sample_dist:
+                t = (next_sample_dist - distance_accum) / edge_len
+                sample = p0 + t * edge_vec
+                samples.append(sample)
+                next_sample_dist += interval
+
+            distance_accum += edge_len
+            edge_index += 1
+
+        if force_last and not np.allclose(samples[-1], points[-1]):
+            samples.append(points[-1])
+
+        return np.array(samples)
+
+    def update(self, data, batch):
+
+        # Get ground truth curves
+        y_curves = batch.y_curve_64.cpu().numpy()  # [total_gt, 64, 3]
+        num_curves_per_batch = batch.num_curves.tolist()
+
+        # Since batch size is 1 in your case
+        B = 1
+
+        # Sample ground truth curves
+        gt_points_list = []
+        gt_cls_list = []
+        for i, gt_curve in enumerate(y_curves):
+            if len(gt_curve) < 2 or np.any(np.isnan(gt_curve)):
+                continue
+            sampled_gt = self.sample_curve_by_interval(
+                gt_curve, self.interval, force_last=True
+            )
+            if len(sampled_gt) > 0 and np.all(np.isfinite(sampled_gt)):
+                gt_points_list.append(sampled_gt)
+                gt_cls_list.append(batch.y_cls[i].cpu().item())
+
+        # Sample predicted curves from post-processed data
+        pred_points_list = []
+        pred_cls_list = []
+        pred_score_list = []
+        for polyline, cls in zip(
+            data.polylines.cpu().numpy(), data.polyline_class.cpu().numpy()
+        ):
+            if cls == 0:  # Skip background class
+                continue
+
+            if len(polyline) < 2 or np.any(np.isnan(polyline)):
+                continue
+
+            sampled_pred = self.sample_curve_by_interval(
+                polyline, self.interval, force_last=True
+            )
+            if len(sampled_pred) > 0 and np.all(np.isfinite(sampled_pred)):
+                pred_points_list.append(sampled_pred)
+                pred_cls_list.append(int(cls))
+                pred_score_list.append(data.polyline_score[i].cpu().item())
+
+        if len(gt_points_list) == 0 and len(pred_points_list) == 0:
+            # No ground truth and no predictions, no penalty
+            self.count += 1
+            return
+        elif len(gt_points_list) == 0:
+            # Penalize no ground truth
+            self.count += 1
+            scores = torch.tensor(pred_score_list)
+            self.map_all_scores.append(scores)
+            self.map_all_matches.append(torch.zeros_like(scores))
+            self.map_all_classes.append(torch.tensor(pred_cls_list))
+            return
+        elif len(pred_points_list) == 0:
+            # Penalize no predictions
+            self.count += 1
+            cls_list = torch.tensor(gt_cls_list, dtype=torch.float32)
+            self.map_all_scores.append(torch.zeros_like(cls_list))
+            self.map_all_matches.append(torch.ones_like(cls_list))
+            self.map_all_classes.append(torch.tensor(cls_list))
+            return
+
+        # calculate mAP
+        for cls in self.map_cls_names.keys():
+            mask = torch.tensor(pred_cls_list) == cls
+            pred_curves = [curve for i, curve in enumerate(pred_points_list) if mask[i]]
+            pred_scores = torch.tensor(pred_score_list)[mask]
+            gt_curves = [
+                curve for i, curve in enumerate(gt_points_list) if cls == gt_cls_list[i]
+            ]
+            if len(pred_curves) == 0 and len(gt_curves) != 0:
+                scores = torch.zeros(len(gt_curves))
+                self.map_all_scores.append(scores)
+                self.map_all_matches.append(torch.zeros_like(scores))
+                self.map_all_classes.append(torch.full_like(scores, cls))
+                continue
+            if len(gt_curves) == 0:
+                self.map_all_scores.append(pred_scores)
+                self.map_all_matches.append(torch.zeros_like(pred_scores))
+                self.map_all_classes.append(torch.full_like(pred_scores, cls))
+                continue
+
+            # get [P, G] matrix of chamfer distances
+            cd_matrix = torch.ones((len(pred_curves), len(gt_curves))) * float("inf")
+            for i, pred_curve in enumerate(pred_curves):
+                for j, gt_curve in enumerate(gt_curves):
+                    cd_matrix[i, j] = calc_chamfer_distance(pred_curve, gt_curve)[0]
+
+            used_gt = set()
+            matches = torch.zeros(len(pred_curves))
+            for i in range(len(pred_curves)):
+                dists = cd_matrix[i]
+                min_dist, min_idx = dists.min(0)
+                if min_dist < self.map_cd_thresh and min_idx not in used_gt:
+                    matches[i] = 1.0
+                    used_gt.add(min_idx)
+
+            self.map_all_scores.append(pred_scores)
+            self.map_all_matches.append(matches)
+            self.map_all_classes.append(torch.full_like(pred_scores, cls))
+
+        pred_points = np.concatenate(pred_points_list, axis=0)
+        gt_points = np.concatenate(gt_points_list, axis=0)
+        # Calculate distances
+        cd, bhd = calc_chamfer_distance(pred_points, gt_points)
+
+        self.total_cd += torch.tensor(cd)
+        self.total_cd_sq += torch.tensor(cd**2)
+        self.total_bhd += torch.tensor(bhd)
+        self.total_bhd_sq += torch.tensor(bhd**2)
+        self.count += 1
+        self.valid_count += 1 if len(pred_points) > 0 else 0
+
+    def compute(self):
+        if not self.map_all_scores:
+            return {cls_name: 0.0 for cls_name in self.map_cls_names.values()}
+        scores = torch.cat(self.map_all_scores)
+        matches = torch.cat(self.map_all_matches)
+        classes = torch.cat(self.map_all_classes)
+        map_result = {}
+        ap_values = []
+
+        for cls in self.map_cls_names.keys():
+            mask = classes == cls
+            if mask.sum() == 0 or torch.sum(matches[mask]) == 0:
+                ap = torch.tensor(0.0, device=self.device)
+            else:
+                ap = average_precision(
+                    scores[mask], matches[mask].to(torch.int32), task="binary"
+                )
+            map_result[self.map_cls_names[cls]] = ap.item()
+            ap_values.append(ap)
+
+        map_result["mAP"] = torch.stack(ap_values).mean().item()
+
+        if self.count == 0:
+            return {"chamfer_distance": 0.0, "bidirectional_hausdorff": 0.0}
+
+        mean_cd = (self.total_cd / self.valid_count).item()
+        mean_bhd = (self.total_bhd / self.valid_count).item()
+        results = {
+            "chamfer_distance": mean_cd,
+            "chamfer_distance_std": (self.total_cd_sq / self.valid_count - (mean_cd**2))
+            .sqrt()
+            .item(),
+            "bidirectional_hausdorff": mean_bhd,
+            "bidirectional_hausdorff_std": (
+                self.total_bhd_sq / self.valid_count - (mean_bhd**2)
+            )
+            .sqrt()
+            .item(),
+        }
+        results.update(map_result)
+        return results

pi3detr/models/__init__.py

@@ -0,0 +1,2 @@
+from .model_config import ModelConfig
+from .pi3detr import PI3DETR

pi3detr/models/losses/__init__.py

@@ -0,0 +1,5 @@
+from .losses import (
+    chamfer_distance_batch,
+    LossParams,
+    ParametricLoss,
+)

pi3detr/models/losses/losses.py

@@ -0,0 +1,399 @@
+from abc import abstractmethod
+import torch
+from torch import nn
+from dataclasses import dataclass, field
+import torch.nn.functional as F
+from torch_geometric.data.data import Data
+from kornia.losses import focal_loss
+from .matcher import *
+
+
+def chamfer_distance_batch(pts1: torch.Tensor, pts2: torch.Tensor) -> torch.Tensor:
+    assert len(pts1.shape) == 3 and len(pts2.shape) == 3
+    if pts1.nelement() == 0 or pts2.nelement() == 0:
+        return torch.tensor(0.0, device=pts1.device, requires_grad=True)
+    dist_matrix = torch.cdist(
+        pts1, pts2, p=2
+    )  # shape: (batch_size, num_points, num_points)
+    dist1 = dist_matrix.min(dim=2).values.mean(dim=1)  # min over pts2, mean over pts1
+    dist2 = dist_matrix.min(dim=1).values.mean(dim=1)  # min over pts1, mean over pts2
+    return (dist1 + dist2) / 2
+
+
+@torch.no_grad()
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    if target.numel() == 0:
+        return [torch.zeros([], device=output.device)]
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+@torch.no_grad()
+def f1_score(output, target, threshold=0.5):
+    output = output > threshold
+    target = target > threshold
+    tp = (output & target).sum()
+    tn = (~output & ~target).sum()
+    fp = (output & ~target).sum()
+    fn = (~output & target).sum()
+    precision = tp / (tp + fp + 1e-8)
+    recall = tp / (tp + fn + 1e-8)
+    return 2 * (precision * recall) / (precision + recall + 1e-8), precision, recall
+
+
+@dataclass
+class LossParams:
+    num_classes: int
+    cost_class: int = 1
+    cost_curve: int = 1
+    class_loss_type: str = "cross_entropy"  # or "focal"
+    class_loss_weights: list[float] = field(
+        default_factory=lambda: [
+            0.04834912,
+            0.40329467,
+            0.09588135,
+            0.23071379,
+            0.22176106,
+        ]
+    )
+    # NOTE: Weights calculated based on the dataset
+    # bezier, line, circle, arc, empty
+    # counts = np.array([11347, 200751, 34672, 37528])
+    # counts = np.append(counts, total_pred - counts.sum())
+    # weights = 1 / np.sqrt(counts)
+    # weights = weights / weights.sum()
+
+
+class Loss(nn.Module):
+
+    def __init__(self, params: LossParams) -> None:
+        super().__init__()
+        self.matcher = ParametricMatcher(params.cost_class, params.cost_curve)
+        self.num_classes = params.num_classes
+        self.class_loss_type = params.class_loss_type
+        class_weights = torch.tensor(
+            params.class_loss_weights,
+        )
+
+        self.register_buffer("class_weights", class_weights)
+
+    def forward(
+        self, outputs: dict[str, torch.Tensor], data: Data
+    ) -> dict[str, torch.Tensor]:
+        indices = self.matcher(outputs, data)
+        losses = {}
+        losses.update(self._loss_class(outputs, data, indices))
+        losses.update(self._loss_polyline(outputs, data, indices))
+        # In case of auxiliary losses, we repeat this process with the output
+        # of each intermediate layer.
+        if "aux_outputs" in outputs:
+            for i, aux_outputs in enumerate(outputs["aux_outputs"]):
+                indices = self.matcher(aux_outputs, data)
+                l_dict = self._loss_class(aux_outputs, data, indices, False)
+                l_dict.update(self._loss_polyline(aux_outputs, data, indices))
+                l_dict = {k + f"_{i}": v for k, v in l_dict.items()}
+                losses.update(l_dict)
+        return losses
+
+    @abstractmethod
+    def _loss_polyline(
+        self,
+        outputs: dict[str, torch.Tensor],
+        data: Data,
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
+    ) -> dict[str, torch.Tensor]:
+        """Compute the polyline loss."""
+        pass
+
+    def _loss_class(
+        self,
+        outputs: dict[str, torch.Tensor],
+        data: Data,
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
+        log: bool = True,
+    ) -> torch.Tensor:
+        num_targets = (
+            data.num_polylines.tolist()
+            if hasattr(data, "num_polylines")
+            else data.num_curves.tolist()
+        )
+        src_logits = outputs["pred_class"]
+        idx = self._get_src_permutation_idx(indices)
+        target_classes_o = torch.cat(
+            [
+                target[J]
+                for target, (_, J) in zip(
+                    data.y_cls.split_with_sizes(num_targets), indices
+                )
+            ]
+        )
+        target_classes = torch.full(
+            src_logits.shape[:2], 0, dtype=torch.int64, device=src_logits.device
+        )  # 0: empty class
+        target_classes[idx] = target_classes_o
+        losses = {}
+
+        if self.class_loss_type == "cross_entropy":
+            loss_class = F.cross_entropy(
+                src_logits.transpose(1, 2),
+                target_classes,
+                weight=self.class_weights.to(src_logits.device),
+                reduction="mean",
+            )
+        else:
+            loss_class = focal_loss(
+                src_logits.transpose(1, 2),
+                target_classes,
+                alpha=0.25,
+                gamma=2.0,
+                weight=self.class_weights.to(src_logits.device),
+                reduction="mean",
+            )
+        losses["loss_class"] = loss_class
+        if log:
+            losses["class_error"] = (
+                100
+                - accuracy(
+                    src_logits.reshape(-1, src_logits.size(-1)),
+                    target_classes.flatten(),
+                )[0]
+            )
+            f1, _, _ = f1_score(
+                src_logits.reshape(-1, src_logits.size(-1)).softmax(-1).argmax(-1),
+                target_classes.flatten(),
+                threshold=0.5,
+            )
+            losses["class_f1_score"] = f1
+
+        return losses
+
+    def _get_src_permutation_idx(
+        self, indices: list[tuple[torch.Tensor, torch.Tensor]]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # permute predictions following indices
+        batch_idx = torch.cat(
+            [torch.full_like(src, i) for i, (src, _) in enumerate(indices)]
+        )
+        src_idx = torch.cat([src for (src, _) in indices])
+        return batch_idx, src_idx
+
+    def _get_tgt_permutation_idx(
+        self, indices: list[tuple[torch.Tensor, torch.Tensor]]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # permute targets following indices
+        batch_idx = torch.cat(
+            [torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)]
+        )
+        tgt_idx = torch.cat([tgt for (_, tgt) in indices])
+        return batch_idx, tgt_idx
+
+
+class ParametricLoss(Loss):
+    def __init__(self, params: LossParams) -> None:
+        super().__init__(params)
+
+    def _loss_polyline(
+        self,
+        outputs: dict[str, torch.Tensor],
+        data: Data,
+        indices: list[tuple[torch.Tensor, torch.Tensor]],
+    ) -> dict[str, torch.Tensor]:
+        idx = self._get_src_permutation_idx(indices)
+        src_bspline_params = outputs["pred_bspline_params"][idx]
+        src_bspline_points = outputs["pred_bspline_points"][idx]
+        src_line_params = outputs["pred_line_params"][idx]
+        src_line_length = outputs["pred_line_length"][idx]
+        src_line_points = outputs["pred_line_points"][idx]
+        src_circle_params = outputs["pred_circle_params"][idx]
+        src_circle_radius = outputs["pred_circle_radius"][idx]
+        src_circle_points = outputs["pred_circle_points"][idx]
+        src_arc_params = outputs["pred_arc_params"][idx]
+        src_arc_points = outputs["pred_arc_points"][idx]
+        target_params = torch.cat(
+            [
+                target[J]
+                for target, (_, J) in zip(
+                    data.y_params.split_with_sizes(data.num_curves.tolist()), indices
+                )
+            ]
+        )
+        target_classes = torch.cat(
+            [
+                target[J]
+                for target, (_, J) in zip(
+                    data.y_cls.split_with_sizes(data.num_curves.tolist()), indices
+                )
+            ]
+        )
+        target_curves = torch.cat(
+            [
+                target[J]
+                for target, (_, J) in zip(
+                    data.y_curve_64.split_with_sizes(data.num_curves.tolist()), indices
+                )
+            ]
+        )
+
+        losses = {}
+
+        # Filter indices for each class
+        bspline_mask = target_classes == 1  # B-spline
+        line_mask = target_classes == 2  # Line
+        circle_mask = target_classes == 3  # Circle
+        arc_mask = target_classes == 4  # Arc
+
+        # Compute loss for B-splines
+        if bspline_mask.any():
+            bspline_order_l1 = torch.min(
+                F.l1_loss(
+                    src_bspline_params[bspline_mask].flatten(-2, -1),
+                    target_params[bspline_mask],
+                    reduction="none",
+                ).mean(-1),
+                F.l1_loss(
+                    src_bspline_params[bspline_mask].flip([1]).flatten(-2, -1),
+                    target_params[bspline_mask],
+                    reduction="none",
+                ).mean(-1),
+            ).mean()
+            losses["loss_bspline"] = bspline_order_l1
+            bspline_chamfer = chamfer_distance_batch(
+                src_bspline_points[bspline_mask], target_curves[bspline_mask]
+            )
+            losses["loss_bspline_chamfer"] = bspline_chamfer.mean()
+        else:
+            losses["loss_bspline"] = torch.tensor(0.0, device=src_bspline_params.device)
+            losses["loss_bspline_chamfer"] = torch.tensor(
+                0.0, device=src_bspline_points.device
+            )
+
+        # Compute loss for Lines
+        if line_mask.any():
+            line_position_l1 = torch.min(
+                F.l1_loss(
+                    src_line_params[line_mask].flatten(-2, -1),
+                    target_params[line_mask, :6],
+                    reduction="none",
+                ).mean(-1),
+                # also consider the negative direction
+                F.l1_loss(
+                    (
+                        src_line_params[line_mask]
+                        * torch.tensor([1.0, -1.0])
+                        .view(1, 2, 1)
+                        .to(src_line_params.device)
+                    ).flatten(-2, -1),
+                    target_params[line_mask, :6],
+                    reduction="none",
+                ).mean(-1),
+            ).mean()
+            line_length_loss = F.l1_loss(
+                src_line_length[line_mask],
+                target_params[line_mask, 6].unsqueeze(-1),
+            )
+            losses["loss_line_position"] = line_position_l1
+            losses["loss_line_length"] = line_length_loss
+            line_chamfer = chamfer_distance_batch(
+                src_line_points[line_mask], target_curves[line_mask]
+            )
+            losses["loss_line_chamfer"] = line_chamfer.mean()
+        else:
+            losses["loss_line_position"] = torch.tensor(
+                0.0, device=src_line_params.device
+            )
+            losses["loss_line_length"] = torch.tensor(
+                0.0, device=src_line_length.device
+            )
+            losses["loss_line_chamfer"] = torch.tensor(
+                0.0, device=src_line_points.device
+            )
+
+        # Compute loss for Circles
+        if circle_mask.any():
+            circle_position_l1 = torch.min(
+                F.l1_loss(
+                    src_circle_params[circle_mask].flatten(-2, -1),
+                    target_params[circle_mask, :6],
+                    reduction="none",
+                ).mean(-1),
+                # also consider the negative direction
+                F.l1_loss(
+                    (
+                        src_circle_params[circle_mask]
+                        * torch.tensor([1.0, -1.0])
+                        .view(1, 2, 1)
+                        .to(src_circle_params.device)
+                    ).flatten(-2, -1),
+                    target_params[circle_mask, :6],
+                    reduction="none",
+                ).mean(-1),
+            ).mean()
+            radius_loss = F.l1_loss(
+                src_circle_radius[circle_mask],
+                target_params[circle_mask, 6].unsqueeze(-1),
+            )
+            losses["loss_circle_position"] = circle_position_l1
+            losses["loss_circle_radius"] = radius_loss
+            circle_chamfer = chamfer_distance_batch(
+                src_circle_points[circle_mask], target_curves[circle_mask]
+            )
+            losses["loss_circle_chamfer"] = circle_chamfer.mean()
+        else:
+            losses["loss_circle_position"] = torch.tensor(
+                0.0, device=src_circle_params.device
+            )
+            losses["loss_circle_radius"] = torch.tensor(
+                0.0, device=src_circle_radius.device
+            )
+            losses["loss_circle_chamfer"] = torch.tensor(
+                0.0, device=src_circle_points.device
+            )
+
+        # Compute loss for Arcs
+        if arc_mask.any():
+            arc_order_l1 = torch.min(
+                F.l1_loss(
+                    src_arc_params[arc_mask].flatten(-2, -1),
+                    target_params[arc_mask, :9],
+                    reduction="none",
+                ).mean(-1),
+                F.l1_loss(
+                    src_arc_params[arc_mask][:, [0, 2, 1]].flatten(-2, -1),
+                    target_params[arc_mask, :9],
+                    reduction="none",
+                ).mean(-1),
+            ).mean()
+            losses["loss_arc"] = arc_order_l1
+            arc_chamfer = chamfer_distance_batch(
+                src_arc_points[arc_mask], target_curves[arc_mask]
+            )
+            losses["loss_arc_chamfer"] = arc_chamfer.mean()
+        else:
+            losses["loss_arc"] = torch.tensor(0.0, device=src_arc_params.device)
+            losses["loss_arc_chamfer"] = torch.tensor(0.0, device=src_arc_points.device)
+
+        losses["total_curve"] = (
+            losses["loss_bspline"]
+            + losses["loss_line_position"]
+            + losses["loss_line_length"]
+            + losses["loss_circle_position"]
+            + losses["loss_circle_radius"]
+            + losses["loss_line_chamfer"]
+            + losses["loss_circle_chamfer"]
+            + losses["loss_bspline_chamfer"]
+            + losses["loss_arc"]
+            + losses["loss_arc_chamfer"]
+        )
+
+        return losses

pi3detr/models/losses/matcher.py

@@ -0,0 +1,135 @@
+import torch
+from torch import nn
+from scipy.optimize import linear_sum_assignment
+from torch_geometric.data.data import Data
+
+
+class ParametricMatcher(nn.Module):
+    def __init__(self, cost_class: int = 1, cost_curve: int = 1) -> None:
+        super().__init__()
+        self.cost_class = cost_class
+        self.cost_curve = cost_curve
+
+    @torch.no_grad()
+    def forward(
+        self, outputs: dict[str, torch.Tensor], data: Data
+    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Compute the matching indices based on class costs and Chamfer distance.
+        """
+        bs, num_queries = outputs["pred_class"].shape[:2]
+
+        # Compute the classification cost
+        out_prob = (
+            outputs["pred_class"].flatten(0, 1).softmax(-1)
+        )  # [batch_size * num_queries, num_classes]
+        cost_class = -out_prob[:, data.y_cls]
+
+        pred_bspline_params = outputs["pred_bspline_params"].flatten(0, 1)
+        pred_line_params = outputs["pred_line_params"].flatten(0, 1)
+        pred_line_length = outputs["pred_line_length"].flatten(0, 1)
+        pred_circle_params = outputs["pred_circle_params"].flatten(0, 1)
+        pred_circle_radius = outputs["pred_circle_radius"].flatten(0, 1)
+        pred_arc_params = outputs["pred_arc_params"].flatten(0, 1)
+
+        # classes -> 1: bspline, 2: line, 3: circle, 4: arc
+        # NOTE: scaling done assuming points are in [-1, 1] range
+        bspline_costs = torch.min(
+            torch.cdist(
+                pred_bspline_params.flatten(-2, -1),
+                data.bspline_params.flatten(-2, -1),
+                p=1,
+            ),
+            torch.cdist(
+                pred_bspline_params.flip([1]).flatten(-2, -1),
+                data.bspline_params.flatten(-2, -1),
+                p=1,
+            ),
+        )  # [batch_size * num_queries, num_curves]
+        line_costs = torch.min(
+            torch.cdist(
+                pred_line_params.flatten(-2, -1),
+                data.line_params.flatten(-2, -1),
+                p=1,
+            ),
+            torch.cdist(
+                (
+                    pred_line_params
+                    * torch.tensor([1.0, -1.0])
+                    .view(1, 2, 1)
+                    .to(pred_line_params.device)
+                ).flatten(-2, -1),
+                data.line_params.flatten(-2, -1),
+                p=1,
+            ),
+        ) + torch.cdist(
+            pred_line_length,
+            data.line_length.unsqueeze(-1),
+            p=1,
+        )  # [batch_size * num_queries, num_curves]
+        circle_costs = torch.min(
+            torch.cdist(
+                pred_circle_params.flatten(-2, -1),
+                data.circle_params.flatten(-2, -1),
+                p=1,
+            ),
+            torch.cdist(
+                (
+                    pred_circle_params
+                    * torch.tensor([1.0, -1.0])
+                    .view(1, 2, 1)
+                    .to(pred_circle_params.device)
+                ).flatten(-2, -1),
+                data.circle_params.flatten(-2, -1),
+                p=1,
+            ),
+        ) + torch.cdist(
+            pred_circle_radius,
+            data.circle_radius.unsqueeze(-1),
+            p=1,
+        )  # [batch_size * num_queries, num_curves]
+        arc_costs = torch.min(
+            torch.cdist(
+                pred_arc_params.flatten(-2, -1),
+                data.arc_params.flatten(-2, -1),
+                p=1,
+            ),
+            # mid, start, end | start and end can be swapped
+            torch.cdist(
+                pred_arc_params[:, [0, 2, 1], :].flatten(-2, -1),
+                data.arc_params.flatten(-2, -1),
+            ),
+        )
+
+        cost_params = torch.stack(
+            [
+                torch.zeros_like(line_costs),
+                bspline_costs,
+                line_costs,
+                circle_costs,
+                arc_costs,
+            ],
+            dim=-1,
+        )
+        cost_params = cost_params[
+            torch.arange(cost_params.size(0))[:, None],
+            torch.arange(cost_params.size(1)),
+            data.y_cls,
+        ]  # [num_queries, num_curves]
+
+        # Combine costs
+        C = self.cost_class * cost_class + self.cost_curve * cost_params
+        C = C.view(bs, num_queries, -1).cpu()
+
+        # Perform Hungarian matching
+        indices = [
+            linear_sum_assignment(c[i])
+            for i, c in enumerate(C.split(data.num_curves.cpu().tolist(), -1))
+        ]
+        return [
+            (
+                torch.as_tensor(i, dtype=torch.int64),
+                torch.as_tensor(j, dtype=torch.int64),
+            )
+            for i, j in indices
+        ]

pi3detr/models/model_config.py

@@ -0,0 +1,66 @@
+from dataclasses import dataclass, field
+from typing import Optional
+from torch_geometric.nn import (
+    MLP,
+)
+from .pointnetpp import SAModule
+
+
+@dataclass
+class ModelConfig:
+    model: str
+    num_features: int
+    epochs: int = 1700
+    lr: float = 1e-4
+    lr_warmup_epochs: int = 15
+    lr_warmup_start_factor: float = 1e-6
+    lr_step: int = 1230
+    batch_size: int = 8
+    batch_size_val: int = 8
+    loss_weights: Optional[dict[str, float]] = None
+    num_curve_points: Optional[int] = 64
+    num_curve_points_val: Optional[int] = 256
+    preencoder_type: Optional[str] = "samodule"
+    preencoder_lr: Optional[float] = 1e-4
+    freeze_backbone: bool = False
+    encoder_dim: Optional[int] = 768
+    decoder_dim: Optional[int] = 768
+    num_encoder_layers: Optional[int] = 3
+    num_decoder_layers: Optional[int] = 9
+    encoder_dropout: float = 0.1
+    decoder_dropout: float = 0.1
+    num_attn_heads: Optional[int] = 8
+    enc_dim_feedforward: Optional[int] = 2048
+    dec_dim_feedforward: Optional[int] = 2048
+    mlp_dropout: float = 0.0
+    num_preds: Optional[int] = 128
+    num_classes: Optional[int] = 5
+    cost_weights: Optional[dict[str, float]] = None
+    auxiliary_loss: bool = True
+    max_points_in_param: Optional[int] = 4
+    num_transformer_points: Optional[int] = 2048
+    query_type: str = "point_fps"
+    pos_embed_type: str = "sine"
+    class_loss_type: str = "cross_entropy"  # or "focal"
+    class_loss_weights: list[float] = field(
+        default_factory=lambda: [
+            0.04834912,
+            0.40329467,
+            0.09588135,
+            0.23071379,
+            0.22176106,
+        ]
+    )
+
+    def get_preencoder(self):
+        preencoder_type = self.preencoder_type
+        preencoder = None
+        if preencoder_type == "samodule":
+            preencoder = SAModule(
+                MLP([self.num_features + 3, 64, 128, self.encoder_dim]),
+                num_out_points=self.num_transformer_points,
+            )
+            preencoder.out_channels = self.encoder_dim
+        else:
+            raise ValueError(f"Unknown preencoder type: {self.preencoder_type}.")
+        return preencoder

pi3detr/models/pi3detr.py

@@ -0,0 +1,593 @@
+import torch
+from torch import nn, Tensor
+import pytorch_lightning as pl
+from pytorch_lightning.utilities.types import OptimizerLRScheduler
+from torch_geometric.data.data import Data
+from torch_geometric.nn import MLP
+from torch_geometric.utils import to_dense_batch
+from .model_config import ModelConfig
+from .losses import LossParams, ParametricLoss
+from .transformer import Transformer
+from .positional_embedding import PositionEmbeddingCoordsSine
+from .query_engine import build_query_engine
+from pi3detr.dataset import reverse_normalize_and_scale
+from ..utils.curve_fitter import (
+    torch_bezier_curve,
+    torch_line_points,
+    generate_points_on_circle_torch,
+    torch_arc_points,
+)
+from ..utils.postprocessing import (
+    snap_and_fit_curves,
+    filter_predictions,
+    iou_filter_point_based,
+    iou_filter_predictions,
+)
+
+from pi3detr.evaluation.abc_metrics import (
+    ChamferMAP,
+)
+from torchmetrics.classification import (
+    BinaryJaccardIndex,
+    BinaryPrecision,
+    BinaryRecall,
+)
+
+
+class PI3DETR(pl.LightningModule):
+    def __init__(self, config: ModelConfig):
+        super().__init__()
+        self.config = config
+        self.pc_preencoder = config.get_preencoder()
+        self.enc_dim = config.encoder_dim
+        self.dec_dim = config.decoder_dim
+        self.num_preds = config.num_preds
+        self.num_curve_points = config.num_curve_points
+        self.num_curve_points_val = config.num_curve_points_val
+        self.num_classes = config.num_classes
+        self.max_points_in_param = config.max_points_in_param
+        self.preenc_to_enc_proj = MLP(
+            [self.pc_preencoder.out_channels, self.enc_dim, self.enc_dim],
+            act="relu",
+            norm="layer_norm",
+        )
+        num_decoder_layers = config.num_decoder_layers
+        self.transformer = Transformer(
+            self.enc_dim,
+            self.dec_dim,
+            nhead=config.num_attn_heads,
+            num_encoder_layers=config.num_encoder_layers,
+            num_decoder_layers=num_decoder_layers,
+            enc_dim_feedforward=config.enc_dim_feedforward,
+            dec_dim_feedforward=config.dec_dim_feedforward,
+            enc_dropout=config.encoder_dropout,
+            dec_dropout=config.decoder_dropout,
+            return_intermediate_dec=True,
+        )
+        self.positional_embedding = PositionEmbeddingCoordsSine(
+            d_pos=self.dec_dim, pos_type=self.config.pos_embed_type
+        )
+        self.pos_embed_proj = MLP(
+            [self.dec_dim, self.dec_dim, self.dec_dim], act="relu", norm="layer_norm"
+        )
+        self.query_type = config.query_type
+        self.query_engine = build_query_engine(
+            self.query_type,
+            self.positional_embedding,
+            self.dec_dim,
+            self.max_points_in_param,
+            self.num_preds,
+        )
+
+        def make_mlp(out_dim, layers=4, base_dim=None, bias_last=True):
+            base_dim = base_dim or self.dec_dim
+            n_layers = layers - 1
+            return MLP(
+                channel_list=[base_dim] * n_layers + [out_dim],
+                bias=[False] * (n_layers - 1) + [bias_last],
+                dropout=self.config.mlp_dropout,
+                act="relu",
+                norm="layer_norm",
+            )
+
+        self.class_head = make_mlp(self.num_classes)
+        self.bspline_param_head = make_mlp(4 * 3)
+        self.line_param_head = make_mlp(2 * 3)
+        self.line_length_head = make_mlp(1, layers=3)
+        self.circle_param_head = make_mlp(2 * 3)
+        self.circle_radius_head = make_mlp(1, layers=3)
+        self.arc_param_head = make_mlp(3 * 3)
+
+        self.loss = ParametricLoss(
+            LossParams(
+                num_classes=self.num_classes - 1,  # -1 for the EOS token
+                cost_class=config.cost_weights["cost_class"],
+                cost_curve=config.cost_weights["cost_curve"],
+                class_loss_type=config.class_loss_type,
+                class_loss_weights=config.class_loss_weights,
+            )
+        )
+        self.auxiliary_loss = self.config.auxiliary_loss
+        self.weight_dict = {
+            "loss_class": config.loss_weights["loss_class"],
+            "loss_bspline": config.loss_weights["loss_bspline"],
+            "loss_bspline_chamfer": config.loss_weights["loss_bspline_chamfer"],
+            "loss_line_position": config.loss_weights["loss_line_position"],
+            "loss_line_length": config.loss_weights["loss_line_length"],
+            "loss_line_chamfer": config.loss_weights["loss_line_chamfer"],
+            "loss_circle_position": config.loss_weights["loss_circle_position"],
+            "loss_circle_radius": config.loss_weights["loss_circle_radius"],
+            "loss_circle_chamfer": config.loss_weights["loss_circle_chamfer"],
+            "loss_arc": config.loss_weights["loss_arc"],
+            "loss_arc_chamfer": config.loss_weights["loss_arc_chamfer"],
+        }
+        # TODO this is a hack
+        self.aux_weight_dict = {}
+        if self.auxiliary_loss:
+            for i in range(num_decoder_layers - 1):
+                self.aux_weight_dict.update(
+                    {k + f"_{i}": v for k, v in self.weight_dict.items()}
+                )
+            self.weight_dict.update(self.aux_weight_dict)
+
+        self.chamfer_map = ChamferMAP(chamfer_thresh=0.05)
+
+        # Torchmetrics for segmentation
+        self.seg_iou = BinaryJaccardIndex()
+        self.seg_precision = BinaryPrecision()
+        self.seg_recall = BinaryRecall()
+
+    def forward(self, data: Data) -> dict[str, Tensor]:
+        x, pos, batch = self.pc_preencoder(data)[-1]
+        x = self.preenc_to_enc_proj(x)
+        x, mask = to_dense_batch(x, batch)
+        pos_dense_batch, _ = to_dense_batch(pos, batch)
+        pos_embed = self.pos_embed_proj(
+            self.positional_embedding(
+                pos_dense_batch, num_channels=self.dec_dim
+            ).permute(0, 2, 1)
+        ).permute(0, 2, 1)
+        query_xyz, query_embed = self.query_engine(Data(pos=pos, batch=batch))
+        x = self.transformer(
+            x,  # [batch_size, num_points, enc_dim]
+            # transformer expects 1s to be masked
+            ~mask if not torch.all(mask) else None,  # [batch_size, num_points]
+            query_embed,  # [batch_size, dec_dim, num_queries]
+            pos_embed,  # [batch_size, dec_dim, num_points]
+        )
+        output_class = self.class_head(x)
+        output_bspline_params = self.bspline_param_head(x)
+        output_line_params = self.line_param_head(x)
+        output_line_length = self.line_length_head(x)
+        output_circle_params = self.circle_param_head(x)
+        output_circle_radius = self.circle_radius_head(x)
+        output_arc_params = self.arc_param_head(x)
+
+        pred_bspline_params = (
+            output_bspline_params[-1].reshape(data.batch_size, self.num_preds, 4, 3)
+            + query_xyz
+        )
+        pred_line_params = output_line_params[-1].reshape(
+            data.batch_size, self.num_preds, 2, 3
+        )
+        pred_line_params[:, :, 0, :] = (
+            pred_line_params[:, :, 0, :] + query_xyz[:, :, 0, :]
+        )
+
+        pred_circle_params = output_circle_params[-1].reshape(
+            data.batch_size, self.num_preds, 2, 3
+        )
+        pred_circle_params[:, :, 0, :] = (
+            pred_circle_params[:, :, 0, :] + query_xyz[:, :, 0, :]
+        )
+
+        pred_arc_params = (
+            output_arc_params[-1].reshape(data.batch_size, self.num_preds, 3, 3)
+            + query_xyz[:, :, :3, :]
+        )
+
+        out = {
+            "pred_class": output_class[-1],
+            "pred_bspline_params": pred_bspline_params,
+            "pred_line_params": pred_line_params,
+            "pred_line_length": output_line_length[-1],
+            "pred_circle_params": pred_circle_params,
+            "pred_circle_radius": output_circle_radius[-1],
+            "pred_arc_params": pred_arc_params,
+            "query_xyz": query_xyz,
+        }
+        if self.auxiliary_loss and self.training:
+            out["aux_outputs"] = self._set_aux_loss(
+                output_bspline_params,
+                output_line_params,
+                output_line_length,
+                output_circle_params,
+                output_circle_radius,
+                output_arc_params,
+                query_xyz,
+                output_class,
+            )
+        return out
+
+    @torch.jit.unused
+    def _set_aux_loss(
+        self,
+        output_bspline_params: torch.Tensor,
+        output_line_params: torch.Tensor,
+        output_line_length: torch.Tensor,
+        output_circle_params: torch.Tensor,
+        output_circle_radius: torch.Tensor,
+        output_arc_params: torch.Tensor,
+        query_xyz: torch.Tensor,
+        output_class: torch.Tensor,
+    ) -> list[dict[str, torch.Tensor]]:
+        # this is a workaround to make torchscript happy, as torchscript
+        # doesn't support dictionary with non-homogeneous values, such
+        # as a dict having both a Tensor and a list.
+        out_aux = []
+        for b, l, ll, c, cr, a, cl in zip(
+            output_bspline_params[:-1],
+            output_line_params[:-1],
+            output_line_length[:-1],
+            output_circle_params[:-1],
+            output_circle_radius[:-1],
+            output_arc_params[:-1],
+            output_class[:-1],
+        ):
+            pred_bspline_params = b.reshape(*b.shape[:2], 4, 3) + query_xyz
+            # second point is the direction vector
+            pred_line_params = l.reshape(*l.shape[:2], 2, 3)
+            pred_line_params_adjusted = pred_line_params.clone()
+            pred_line_params_adjusted[:, :, 0, :] = (
+                pred_line_params[:, :, 0, :] + query_xyz[:, :, 0, :]
+            )
+            pred_circle_params = c.reshape(*c.shape[:2], 2, 3)
+            pred_circle_params_adjusted = pred_circle_params.clone()
+            pred_circle_params_adjusted[:, :, 0, :] = (
+                pred_circle_params[:, :, 0, :] + query_xyz[:, :, 0, :]
+            )
+            pred_arc_params = a.reshape(*a.shape[:2], 3, 3) + query_xyz[:, :, :3, :]
+
+            layer_out = {
+                "pred_bspline_params": pred_bspline_params,
+                "pred_line_params": pred_line_params_adjusted,
+                "pred_line_length": ll,
+                "pred_circle_params": pred_circle_params_adjusted,
+                "pred_circle_radius": cr,
+                "pred_arc_params": pred_arc_params,
+                "pred_class": cl,
+            }
+            layer_out.update(
+                self._sample_curve_points(layer_out, self.num_curve_points)
+            )
+            out_aux.append(layer_out)
+
+        return out_aux
+
+    def _sample_curve_points(
+        self, out: dict[str, Tensor], num_points: int
+    ) -> dict[str, Tensor]:
+        batch_size, num_preds = out["pred_bspline_params"].shape[:2]
+        pred_line_params = out["pred_line_params"]
+        pred_line_length = out["pred_line_length"]
+        pred_line_start = (
+            pred_line_params[:, :, 0, :]
+            - pred_line_params[:, :, 1, :] * pred_line_length / 2.0
+        )
+        pred_line_end = (
+            pred_line_params[:, :, 0, :]
+            + pred_line_params[:, :, 1, :] * pred_line_length / 2.0
+        )
+        curves = {}
+        curves["pred_bspline_points"] = torch_bezier_curve(
+            out["pred_bspline_params"].reshape(-1, 4, 3), num_points
+        ).reshape(batch_size, num_preds, -1, 3)
+        curves["pred_line_points"] = torch_line_points(
+            pred_line_start.reshape(-1, 3),
+            pred_line_end.reshape(-1, 3),
+            num_points,
+        ).reshape(batch_size, num_preds, -1, 3)
+        curves["pred_circle_points"] = generate_points_on_circle_torch(
+            out["pred_circle_params"].reshape(-1, 2, 3)[:, 0],
+            out["pred_circle_params"].reshape(-1, 2, 3)[:, 1],
+            out["pred_circle_radius"].reshape(-1),
+            num_points,
+        ).reshape(batch_size, num_preds, -1, 3)
+        curves["pred_arc_points"] = torch_arc_points(
+            out["pred_arc_params"][:, :, 1, :].reshape(-1, 3),
+            out["pred_arc_params"][:, :, 0, :].reshape(-1, 3),
+            out["pred_arc_params"][:, :, 2, :].reshape(-1, 3),
+            num_points,
+        ).reshape(batch_size, num_preds, -1, 3)
+        return curves
+
+    def predict_step(
+        self,
+        batch: Data,
+        reverse_norm: bool = True,
+        thresholds: list[float] = None,
+        snap_and_fit: bool = True,
+        iou_filter: bool = False,
+    ) -> list[Data]:
+        preds = self(batch)
+        preds.update(self._sample_curve_points(preds, self.num_curve_points_val))
+
+        outputs = self.decode_predictions(batch, preds, reverse_norm)
+
+        if thresholds:
+            outputs = [filter_predictions(data, thresholds) for data in outputs]
+
+        if snap_and_fit:
+            outputs = [snap_and_fit_curves(data.clone()) for data in outputs]
+
+        if iou_filter:
+            outputs = [iou_filter_predictions(data) for data in outputs]
+
+        return outputs
+
+    def training_step(self, batch: Data, batch_idx: int) -> Tensor:
+        outputs = self(batch)
+        outputs.update(self._sample_curve_points(outputs, self.num_curve_points))
+        loss_dict = self.loss(outputs, batch)
+        for k, v in loss_dict.items():
+            weight = self.weight_dict[k] if "loss" in k else 1
+            self._default_log(k, v * weight)
+        # weigh losses and sum them for backpropagation
+        weighted_loss_dict = {
+            k: loss_dict[k] * self.weight_dict[k] for k in self.weight_dict.keys()
+        }
+        total_loss = sum(weighted_loss_dict.values())
+        self._default_log("loss_train", total_loss)
+        return total_loss
+
+    @torch.no_grad()
+    def validation_step(self, batch: Data, batch_idx: int) -> None:
+        outputs = self(batch)
+        # sample the training curve points for loss computation
+        outputs.update(self._sample_curve_points(outputs, self.num_curve_points))
+        loss_dict = self.loss(outputs, batch)
+        for k, v in loss_dict.items():
+            weight = self.weight_dict[k] if "loss" in k else 1
+            self._default_log(f"val_{k}", v * weight)
+        without_aux = {
+            k for k in self.weight_dict.keys() if k not in self.aux_weight_dict
+        }
+        self._default_log(
+            "loss_val", sum(loss_dict[k] * self.weight_dict[k] for k in without_aux)
+        )
+        # Sample curve points for validation
+        outputs.update(self._sample_curve_points(outputs, self.num_curve_points_val))
+        self._compute_metrics(batch, outputs)
+
+    @torch.no_grad()
+    def on_validation_epoch_end(self) -> None:
+        metrics = self.chamfer_map.compute()
+        for key, value in metrics.items():
+            self._default_log(f"val_{key}", value)
+        self.chamfer_map.reset()
+
+        # Log segmentation metrics at epoch end
+        self._default_log("val_seg_iou", self.seg_iou.compute())
+        self._default_log(
+            "val_seg_precision",
+            self.seg_precision.compute(),
+        )
+        self._default_log("val_seg_recall", self.seg_recall.compute())
+        self.seg_iou.reset()
+        self.seg_precision.reset()
+        self.seg_recall.reset()
+
+    def test_step(self, batch: Data, batch_idx: int) -> None:
+        outputs = self(batch)
+        outputs.update(self._sample_curve_points(outputs, self.num_curve_points_val))
+        self._compute_metrics(batch, outputs)
+
+    def on_test_epoch_end(self) -> None:
+        metrics = self.chamfer_map.compute()
+        self.chamfer_map.reset()
+        for key, value in metrics.items():
+            self.log(f"test_{key}", value, prog_bar=False)
+
+        # Log segmentation metrics at epoch end
+        self.log("test_seg_iou", self.seg_iou.compute(), prog_bar=False)
+        self.log("test_seg_precision", self.seg_precision.compute(), prog_bar=False)
+        self.log("test_seg_recall", self.seg_recall.compute(), prog_bar=False)
+        self.seg_iou.reset()
+        self.seg_precision.reset()
+        self.seg_recall.reset()
+
+    def _compute_metrics(self, batch: Data, preds: dict):
+        # segmentation metrics
+        outputs = self.decode_predictions(batch, preds, reverse_norm=True)
+        for i, output in enumerate(outputs):
+            self.seg_iou.update(output.segmentation, output.y_seg)
+            self.seg_precision.update(output.segmentation, output.y_seg)
+            self.seg_recall.update(output.segmentation, output.y_seg)
+        # chamfer metrics
+        self.chamfer_map.update(preds, batch)
+
+    def set_num_preds(self, num_preds: int) -> None:
+        if num_preds == self.num_preds:
+            return
+        self.num_preds = num_preds
+        old_state = (
+            self.query_engine.state_dict()
+            if isinstance(self.query_engine, nn.Module)
+            else None
+        )
+        new_engine = build_query_engine(
+            self.query_type,
+            self.positional_embedding,
+            self.dec_dim,
+            self.max_points_in_param,
+            self.num_preds,
+        )
+        if old_state is not None:
+            new_state = new_engine.state_dict()
+            for k, v in old_state.items():
+                assert k in new_state, f"Missing parameter in new query engine: {k}"
+                nv = new_state[k]
+                assert (
+                    v.shape == nv.shape
+                ), f"Shape mismatch for {k}: {v.shape} != {nv.shape}"
+                nv.copy_(v.to(nv.device))
+            new_engine.load_state_dict(new_state, strict=True)
+        self.query_engine = new_engine.to(self.device)
+
+    @torch.no_grad()
+    def decode_predictions(
+        self, batch: Data, preds: Data, reverse_norm: bool = True
+    ) -> list[Data]:
+        outputs = []
+
+        # Vectorized class prediction and score
+        preds_class = preds["pred_class"].softmax(-1)
+        polyline_class = preds_class.argmax(-1)  # (batch_size, num_preds)
+        polyline_score = preds_class.max(-1).values  # (batch_size, num_preds)
+
+        # Prepare all possible polylines: (batch_size, num_preds, num_polypoints, 3)
+        bspline_points = preds["pred_bspline_points"]
+        line_points = preds["pred_line_points"]
+        circle_points = preds["pred_circle_points"]
+        arc_points = preds["pred_arc_points"]
+        zeros_points = torch.zeros_like(bspline_points)  # EOS/empty
+
+        # Stack all types: (batch_size, num_preds, 4, num_polypoints, 3)
+        all_polylines = torch.stack(
+            [zeros_points, bspline_points, line_points, circle_points, arc_points],
+            dim=2,
+        )
+
+        # Gather correct polyline for each prediction
+        # polyline_class: (batch_size, num_preds)
+        # Need to expand to match all_polylines shape for gather
+        idx = (
+            polyline_class.unsqueeze(-1)
+            .unsqueeze(-1)
+            .unsqueeze(2)  # shape: (batch_size, num_preds, 1, num_polypoints, 3)
+            .expand(-1, -1, 1, self.num_curve_points_val, 3)
+        )
+        polylines = torch.gather(all_polylines, 2, idx).squeeze(
+            2
+        )  # (batch_size, num_preds, num_polypoints, 3)
+
+        batch_size = batch.batch_size
+        device = batch.pos.device
+        segmentations = []
+        for i in range(batch_size):
+            # If all predicted classes are zero (EOS), segmentation should be all zeros
+            if torch.all(polyline_class[i] == 0):
+                pc_pts = batch.pos[batch.batch == i]
+                segmentation = torch.zeros(
+                    pc_pts.shape[0], dtype=torch.long, device=device
+                )
+                segmentations.append(segmentation)
+                continue
+
+            poly_pts = polylines[i, polyline_class[i] != 0].reshape(-1, 3)
+            pc_pts = batch.pos[batch.batch == i]  # (num_points_in_cloud, 3)
+            dists = torch.cdist(poly_pts, pc_pts)
+            closest_idx = dists.argmin(dim=1)
+            segmentation = torch.zeros(pc_pts.shape[0], dtype=torch.long, device=device)
+            segmentation[closest_idx.unique()] = 1
+            segmentations.append(segmentation)
+
+        for i in range(batch.batch_size):
+            output = Data(
+                pos=batch.pos[batch.batch == i].clone(),  # point cloud
+                bspline_points=bspline_points[i],  # prediction of B-spline head
+                line_points=line_points[i],  # prediction of line heads
+                circle_points=circle_points[i],  # prediction of circle heads
+                arc_points=arc_points[i],  # prediction of arc head
+                polyline_class=polyline_class[i],  # class of each polyline
+                polyline_score=polyline_score[i],  # score of polyline class
+                polylines=polylines[i],  # polyline that matches polyline_class
+                segmentation=segmentations[
+                    i
+                ],  # curve segmentation for whole point cloud
+                query_xyz=preds["query_xyz"][i],  # query points for the transformer
+            )
+            if hasattr(batch, "y_seg"):
+                output.y_seg = batch.y_seg[batch.batch == i]
+
+            if reverse_norm:
+                output.center = batch.center[i]
+                output.scale = batch.scale[i]
+
+                output = reverse_normalize_and_scale(
+                    output,
+                    extra_fields=[
+                        "polylines",
+                        "bspline_points",
+                        "line_points",
+                        "circle_points",
+                        "arc_points",
+                        "query_xyz",
+                    ],
+                )
+            outputs.append(output)
+
+        return outputs
+
+    def _default_log(self, name: str, value: Tensor) -> None:
+        batch_size = (
+            self.config.batch_size if self.training else self.config.batch_size_val
+        )
+        self.log(
+            name,
+            value,
+            prog_bar=True,
+            on_epoch=True,
+            on_step=False,
+            sync_dist=True,
+            batch_size=batch_size,
+        )
+
+    def configure_optimizers(self) -> OptimizerLRScheduler:
+        param_dict = None
+        config = self.config
+        if config.freeze_backbone:
+            for param in self.pc_preencoder.parameters():
+                param.requires_grad = False
+            param_dict = self.parameters()
+        elif config.lr != config.preencoder_lr:
+            param_dict = [
+                {
+                    "params": [
+                        p for n, p in self.named_parameters() if "pc_encoder" not in n
+                    ]
+                },
+                {
+                    "params": [
+                        p for n, p in self.named_parameters() if "pc_encoder" in n
+                    ],
+                    "lr": self.config.preencoder_lr,
+                },
+            ]
+        else:
+            param_dict = self.parameters()
+        # ----- OPTIMIZER -----
+        optimizer = torch.optim.AdamW(param_dict, lr=self.config.lr)
+        # ----- WARMUP SCHEDULER -----
+        warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+            optimizer,
+            start_factor=config.lr_warmup_start_factor,  # start near zero
+            end_factor=1.0,  # ramp up to base LR
+            total_iters=config.lr_warmup_epochs,
+        )
+        # ----- STEP SCHEDULER -----
+        # Drop LR by factor after (step_epoch - warmup_epochs) epochs
+        step_scheduler = torch.optim.lr_scheduler.StepLR(
+            optimizer,
+            step_size=config.lr_step - config.lr_warmup_epochs,
+            gamma=0.1,  # drop LR to 10%
+            last_epoch=config.epochs,
+        )
+        # ----- COMBINE -----
+        scheduler = torch.optim.lr_scheduler.SequentialLR(
+            optimizer,
+            schedulers=[warmup_scheduler, step_scheduler],
+            milestones=[config.lr_warmup_epochs],
+        )
+
+        return [optimizer], {"scheduler": scheduler, "interval": "epoch"}

pi3detr/models/pointnetpp.py

@@ -0,0 +1,324 @@
+import torch
+from torch import nn, Tensor
+from typing import Optional, Tuple
+import numpy as np
+from scipy.spatial import cKDTree
+from torch_geometric.nn import (
+    MLP,
+    PointNetConv,
+    fps,
+    knn_interpolate,
+    radius,
+    global_max_pool,
+)
+from torch_geometric.data.data import BaseData
+
+TensorTriple = Tuple[Tensor, Tensor, Tensor]
+
+
+def radius_cpu(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    r: float,
+    batch_x: Optional[torch.Tensor] = None,
+    batch_y: Optional[torch.Tensor] = None,
+    max_num_neighbors: Optional[int] = None,
+    loop: bool = False,
+    sort_by_distance: bool = True,
+) -> Tuple[torch.LongTensor, torch.LongTensor]:
+    """
+    CPU replacement for torch_cluster.radius / torch_geometric.radius.
+
+    Semantics (matching torch_geometric.radius):
+      Returns (row, col) where
+        row: indices into `y` (centers) in the range [0, y.size(0))
+        col: indices into `x` (neighbors) in the range [0, x.size(0))
+
+    Thus, for y = x[idx]:
+      edge_index = torch.stack([col, row], dim=0)
+      edge_index[0] indexes the full set (source/neighbor),
+      edge_index[1] indexes the sampled centers.
+    """
+    # Basic checks
+    if x.device.type != "cpu" or y.device.type != "cpu":
+        raise ValueError("radius_cpu expects x and y to be on CPU.")
+    if x.ndim != 2 or y.ndim != 2:
+        raise ValueError("x and y must be 2D (N, D).")
+    if x.shape[1] != y.shape[1]:
+        raise ValueError("x and y must have same dimensionality D.")
+
+    N_x = x.shape[0]
+    N_y = y.shape[0]
+    if N_x == 0 or N_y == 0:
+        return torch.empty((0,), dtype=torch.long), torch.empty((0,), dtype=torch.long)
+
+    x_np = np.asarray(x)
+    y_np = np.asarray(y)
+
+    if batch_x is None:
+        batch_x = torch.zeros(N_x, dtype=torch.long)
+    else:
+        if batch_x.device.type != "cpu":
+            batch_x = batch_x.cpu()
+        batch_x = batch_x.long()
+
+    if batch_y is None:
+        batch_y = torch.zeros(N_y, dtype=torch.long)
+    else:
+        if batch_y.device.type != "cpu":
+            batch_y = batch_y.cpu()
+        batch_y = batch_y.long()
+
+    rows = []
+    cols = []
+
+    unique_batches = torch.unique(torch.cat([batch_x, batch_y])).tolist()
+    # iterate only over batches actually present in y to avoid unnecessary work
+    unique_batches = sorted(set(batch_y.tolist()))
+
+    for b in unique_batches:
+        # mask and maps from local->global indices
+        mask_x = (batch_x == b).numpy()
+        mask_y = (batch_y == b).numpy()
+        idxs_x = np.nonzero(mask_x)[0]  # global indices in x
+        idxs_y = np.nonzero(mask_y)[0]  # global indices in y
+
+        if idxs_y.size == 0 or idxs_x.size == 0:
+            continue
+
+        pts_x = x_np[mask_x]
+        pts_y = y_np[mask_y]
+
+        # build tree on source points (x) and query for each center in y
+        tree = cKDTree(pts_x)
+        # neighbors_list: for each center (local), a list of local indices into pts_x
+        neighbors_list = tree.query_ball_point(pts_y, r)
+
+        for local_center, neigh_locals in enumerate(neighbors_list):
+            if len(neigh_locals) == 0:
+                continue
+            neigh_locals = np.array(neigh_locals, dtype=int)
+
+            # remove self if requested AND x and y are the same set at same global indices
+            if not loop:
+                # If x and y refer to the same global indices and same coords, remove self-match
+                # we detect self by checking whether global index equals center global index
+                center_global = idxs_y[local_center]
+                # compute global neighbor indices
+                neigh_globals = idxs_x[neigh_locals]
+                # boolean mask for neighbors that are not self
+                not_self_mask = neigh_globals != center_global
+                neigh_locals = neigh_locals[not_self_mask]
+
+                if neigh_locals.size == 0:
+                    continue
+
+            # apply max_num_neighbors: keep closest ones by distance if requested
+            if max_num_neighbors is not None and neigh_locals.size > max_num_neighbors:
+                if sort_by_distance:
+                    dists = np.linalg.norm(
+                        pts_x[neigh_locals] - pts_y[local_center], axis=1
+                    )
+                    order = np.argsort(dists)[:max_num_neighbors]
+                    neigh_locals = neigh_locals[order]
+                else:
+                    neigh_locals = np.sort(neigh_locals)[:max_num_neighbors]
+
+            # optionally sort by distance
+            if sort_by_distance and neigh_locals.size > 0:
+                dists = np.linalg.norm(
+                    pts_x[neigh_locals] - pts_y[local_center], axis=1
+                )
+                order = np.argsort(dists)
+                neigh_locals = neigh_locals[order]
+
+            # convert to global indices and append
+            neigh_globals = idxs_x[neigh_locals].tolist()
+            center_global = int(idxs_y[local_center])
+            rows.extend(neigh_globals)  # neighbor indices into x (row)
+            cols.extend(
+                [center_global] * len(neigh_globals)
+            )  # center indices into y (col)
+
+    if len(rows) == 0:
+        return torch.empty((0,), dtype=torch.long), torch.empty((0,), dtype=torch.long)
+
+    row_t = torch.tensor(rows, dtype=torch.long)  # currently neighbors (x)
+    col_t = torch.tensor(cols, dtype=torch.long)  # currently centers (y)
+
+    # Swap to enforce (row=center_indices_in_y, col=neighbor_indices_in_x)
+    return col_t, row_t
+
+
+class SAModuleRatio(torch.nn.Module):
+    def __init__(
+        self, ratio: float, r: float, nn: nn.Module, max_num_neighbors: int = 64
+    ):
+        super().__init__()
+        self.ratio = ratio
+        self.r = r
+        self.conv = PointNetConv(nn, add_self_loops=False)
+        self.max_num_neighbors = max_num_neighbors
+
+    def forward(self, x: torch.Tensor, pos: torch.Tensor, batch: torch.Tensor):
+        idx = fps(pos, batch, ratio=self.ratio)
+        row, col = radius(
+            pos,
+            pos[idx],
+            self.r,
+            batch,
+            batch[idx],
+            max_num_neighbors=self.max_num_neighbors,
+        )
+        edge_index = torch.stack([col, row], dim=0)
+        x_dst = None if x is None else x[idx]
+        x = self.conv((x, x_dst), (pos, pos[idx]), edge_index)
+        pos, batch = pos[idx], batch[idx]
+        return x, pos, batch
+
+
+class SAModule(torch.nn.Module):
+    def __init__(
+        self,
+        nn: nn.Module,
+        num_out_points: float = 2048,
+        r: float = 0.2,
+        max_num_neighbors: int = 64,
+    ):
+        super().__init__()
+        self.num_out_points = num_out_points
+        self.r = r
+        self.conv = PointNetConv(nn, add_self_loops=False)
+        self.max_num_neighbors = max_num_neighbors
+
+    def forward(self, data: BaseData) -> list[tuple[TensorTriple]]:
+        x, pos, batch = data.x, data.pos, data.batch
+        num_points_per_batch = torch.bincount(batch)
+        max_ratio = self.num_out_points / num_points_per_batch.min().item()
+        fps_idx = fps(pos, batch, ratio=max_ratio)
+        fps_batch = batch[fps_idx]
+        idx = torch.cat(
+            [
+                fps_idx[fps_batch == i][: self.num_out_points]
+                for i in range(batch.max().item() + 1)
+            ]
+        )
+        if pos.device == torch.device("cpu"):
+            row, col = radius_cpu(
+                pos,
+                pos[idx],
+                self.r,
+                batch,
+                batch[idx],
+                max_num_neighbors=self.max_num_neighbors,
+                sort_by_distance=False,
+            )
+        else:  # GPU
+            row, col = radius(
+                pos,
+                pos[idx],
+                self.r,
+                batch,
+                batch[idx],
+                max_num_neighbors=self.max_num_neighbors,
+            )
+        edge_index = torch.stack([col, row], dim=0)
+        x_dst = None if x is None else x[idx]
+        x = self.conv((x, x_dst), (pos, pos[idx]), edge_index)
+        pos, batch = pos[idx], batch[idx]
+        return [(x, pos, batch)]
+
+
+class GlobalSAModule(torch.nn.Module):
+    def __init__(self, nn):
+        super().__init__()
+        self.nn = nn
+
+    def forward(self, x: torch.Tensor, pos: torch.Tensor, batch: torch.Tensor):
+        x = self.nn(torch.cat([x, pos], dim=1))
+        x = global_max_pool(x, batch)
+        pos = pos.new_zeros((x.size(0), 3))
+        batch = torch.arange(x.size(0), device=batch.device)
+        return x, pos, batch
+
+
+class FPModule(nn.Module):
+    def __init__(self, k, nn):
+        super().__init__()
+        self.k = k
+        self.nn = nn
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        pos: torch.Tensor,
+        batch: torch.Tensor,
+        x_skip: torch.Tensor,
+        pos_skip: torch.Tensor,
+        batch_skip: torch.Tensor,
+    ):
+        x = knn_interpolate(x, pos, pos_skip, batch, batch_skip, k=self.k)
+        if x_skip is not None:
+            x = torch.cat([x, x_skip], dim=1)
+        x = self.nn(x)
+        return x, pos_skip, batch_skip
+
+
+class PointNetPPEncoder(nn.Module):
+    def __init__(self, num_features: int = 3, out_channels: int = 512):
+        super().__init__()
+        self.out_channels = out_channels
+        # Input channels account for both `pos` and node features.
+        self.sa1_module = SAModuleRatio(
+            0.5, 0.05, MLP([num_features + 3, 32, 32, 64]), 32
+        )
+        self.sa2_module = SAModuleRatio(0.5, 0.1, MLP([64 + 3, 64, 64, 128]), 32)
+        self.sa3_module = SAModuleRatio(0.5, 0.2, MLP([128 + 3, 128, 128, 256]), 32)
+        self.sa4_module = SAModuleRatio(
+            0.5, 0.4, MLP([256 + 3, 256, 256, self.out_channels]), 32
+        )
+
+    def forward(self, data: BaseData) -> list[Tensor]:
+        sa0_out = (data.x, data.pos, data.batch)
+        sa1_out = self.sa1_module(*sa0_out)
+        sa2_out = self.sa2_module(*sa1_out)
+        sa3_out = self.sa3_module(*sa2_out)
+        sa4_out = self.sa4_module(*sa3_out)
+        return [sa0_out, sa1_out, sa2_out, sa3_out, sa4_out]
+
+
+class PointNetPPDecoder(nn.Module):
+    def __init__(self, num_features: int = 3, out_channels: int = 256):
+        super().__init__()
+        self.out_channels = out_channels
+        self.fp4_module = FPModule(1, MLP([512 + 256, 256, 256]))
+        self.fp3_module = FPModule(3, MLP([256 + 128, 256, 256]))
+        self.fp2_module = FPModule(3, MLP([256 + 64, 256, 128]))
+        self.fp1_module = FPModule(3, MLP([128 + num_features, 128, self.out_channels]))
+
+    def forward(
+        self,
+        sa0_out: TensorTriple,
+        sa1_out: TensorTriple,
+        sa2_out: TensorTriple,
+        sa3_out: TensorTriple,
+        sa4_out: TensorTriple,
+    ) -> TensorTriple:
+        fp4_out = self.fp4_module(*sa4_out, *sa3_out)
+        fp3_out = self.fp3_module(*fp4_out, *sa2_out)
+        fp2_out = self.fp2_module(*fp3_out, *sa1_out)
+        x, pos, batch = self.fp1_module(*fp2_out, *sa0_out)
+        return [x, pos, batch]
+
+
+class PointNetPP(nn.Module):
+    def __init__(self, num_features: int, dec_out_channels: int = 256):
+        super().__init__()
+        self.encoder = PointNetPPEncoder(num_features)
+        self.decoder = PointNetPPDecoder(num_features, dec_out_channels)
+        self.out_channels = self.decoder.out_channels
+
+    def forward(self, data: BaseData) -> TensorTriple:
+        x = self.encoder(data)
+        x, pos, batch = self.decoder(*x)
+        return [(x, pos, batch)]

pi3detr/models/positional_embedding.py

@@ -0,0 +1,133 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+
+
+
+
+"""
+Various positional encodings for the transformer.
+"""
+import math
+import torch
+from torch import nn
+import numpy as np
+
+class PositionEmbeddingCoordsSine(nn.Module):
+    def __init__(
+        self,
+        temperature=10000,
+        normalize=False,
+        scale=None,
+        pos_type="fourier",
+        d_pos=None,
+        d_in=3,
+        gauss_scale=1.0,
+    ):
+        super().__init__()
+        self.temperature = temperature
+        self.normalize = normalize
+        if scale is not None and normalize is False:
+            raise ValueError("normalize should be True if scale is passed")
+        if scale is None:
+            scale = 2 * math.pi
+        assert pos_type in ["sine", "fourier"]
+        self.pos_type = pos_type
+        self.scale = scale
+        if pos_type == "fourier":
+            assert d_pos is not None
+            assert d_pos % 2 == 0
+            # define a gaussian matrix input_ch -> output_ch
+            B = torch.empty((d_in, d_pos // 2)).normal_()
+            B *= gauss_scale
+            self.register_buffer("gauss_B", B)
+            self.d_pos = d_pos
+
+    def get_sine_embeddings(self, xyz, num_channels):
+        # clone coords so that shift/scale operations do not affect original tensor
+        orig_xyz = xyz
+        xyz = orig_xyz.clone()
+
+        ndim = num_channels // xyz.shape[2]
+        if ndim % 2 != 0:
+            ndim -= 1
+        # automatically handle remainder by assiging it to the first dim
+        rems = num_channels - (ndim * xyz.shape[2])
+
+        assert (
+            ndim % 2 == 0
+        ), f"Cannot handle odd sized ndim={ndim} where num_channels={num_channels} and xyz={xyz.shape}"
+
+        final_embeds = []
+        prev_dim = 0
+
+        for d in range(xyz.shape[2]):
+            cdim = ndim
+            if rems > 0:
+                # add remainder in increments of two to maintain even size
+                cdim += 2
+                rems -= 2
+
+            if cdim != prev_dim:
+                dim_t = torch.arange(cdim, dtype=torch.float32, device=xyz.device)
+                dim_t = self.temperature ** (2 * (dim_t // 2) / cdim)
+
+            # create batch x cdim x nccords embedding
+            raw_pos = xyz[:, :, d]
+            if self.scale:
+                raw_pos *= self.scale
+            pos = raw_pos[:, :, None] / dim_t
+            pos = torch.stack(
+                (pos[:, :, 0::2].sin(), pos[:, :, 1::2].cos()), dim=3
+            ).flatten(2)
+            final_embeds.append(pos)
+            prev_dim = cdim
+
+        final_embeds = torch.cat(final_embeds, dim=2).permute(0, 2, 1)
+        return final_embeds
+
+    def get_fourier_embeddings(self, xyz, num_channels=None):
+        # Follows - https://people.eecs.berkeley.edu/~bmild/fourfeat/index.html
+
+        if num_channels is None:
+            num_channels = self.gauss_B.shape[1] * 2
+
+        bsize, npoints = xyz.shape[0], xyz.shape[1]
+        assert num_channels > 0 and num_channels % 2 == 0
+        d_in, max_d_out = self.gauss_B.shape[0], self.gauss_B.shape[1]
+        d_out = num_channels // 2
+        assert d_out <= max_d_out
+        assert d_in == xyz.shape[-1]
+
+        # clone coords so that shift/scale operations do not affect original tensor
+        orig_xyz = xyz
+        xyz = orig_xyz.clone()
+
+        xyz *= 2 * np.pi
+        xyz_proj = torch.mm(xyz.view(-1, d_in), self.gauss_B[:, :d_out]).view(
+            bsize, npoints, d_out
+        )
+        final_embeds = [xyz_proj.sin(), xyz_proj.cos()]
+
+        # return batch x d_pos x npoints embedding
+        final_embeds = torch.cat(final_embeds, dim=2).permute(0, 2, 1)
+        return final_embeds
+
+    def forward(self, xyz, num_channels=None):
+        assert isinstance(xyz, torch.Tensor)
+        assert xyz.ndim == 3
+        # xyz is batch x npoints x 3
+        if self.pos_type == "sine":
+            with torch.no_grad():
+                return self.get_sine_embeddings(xyz, num_channels)
+        elif self.pos_type == "fourier":
+            with torch.no_grad():
+                return self.get_fourier_embeddings(xyz, num_channels)
+        else:
+            raise ValueError(f"Unknown {self.pos_type}")
+
+    def extra_repr(self):
+        st = f"type={self.pos_type}, scale={self.scale}, normalize={self.normalize}"
+        if hasattr(self, "gauss_B"):
+            st += (
+                f", gaussB={self.gauss_B.shape}, gaussBsum={self.gauss_B.sum().item()}"
+            )
+        return st

pi3detr/models/query_engine.py

@@ -0,0 +1,108 @@
+import torch
+from torch import nn, Tensor
+from abc import ABC, abstractmethod
+from torch_geometric.nn import MLP, fps, knn
+from torch_geometric.data.data import Data
+from typing import Optional
+
+
+class QueryEngine(nn.Module, ABC):
+    def __init__(
+        self,
+        pos_embedder: Optional[nn.Module],
+        feat_dim: int,
+        max_points_in_param: int,
+        num_queries: int,
+    ):
+        super().__init__()
+        self.pos_embedder = pos_embedder
+        self.feat_dim = feat_dim
+        self.max_points_in_param = max_points_in_param
+        self.num_queries = num_queries
+
+    @abstractmethod
+    def forward(self, data: Data) -> tuple[Tensor]:
+        pass
+
+
+class PointFPSQueryEngine(QueryEngine):
+    def __init__(
+        self,
+        pos_embedder: nn.Module,
+        feat_dim: int,
+        max_points_in_param: int,
+        num_queries: int,
+    ):
+        super().__init__(pos_embedder, feat_dim, max_points_in_param, num_queries)
+        self.num_queries = num_queries
+        self.query_proj = MLP(
+            [self.feat_dim, self.feat_dim, self.feat_dim],
+            bias=False,
+            act="relu",
+            norm="layer_norm",
+        )
+
+    def forward(self, data: Data) -> tuple[Tensor]:
+        num_points_per_batch = torch.bincount(data.batch)
+        max_ratio = self.num_queries / num_points_per_batch.min().item()
+        fps_idx = fps(data.pos, data.batch, ratio=max_ratio)
+        fps_batch = data.batch[fps_idx]
+        query_xyz = torch.stack(
+            [
+                data.pos[fps_idx[fps_batch == i][: self.num_queries]]
+                for i in range(data.batch.max().item() + 1)
+            ]
+        )
+        query_pos = self.pos_embedder(query_xyz, num_channels=self.feat_dim)
+        query_embed = self.query_proj(query_pos.permute(0, 2, 1))[
+            :, : self.num_queries, :
+        ].permute(0, 2, 1)
+        return (
+            query_xyz.unsqueeze(2).expand(-1, -1, self.max_points_in_param, -1),
+            query_embed,
+        )
+
+
+class LearnedQueryEngine(QueryEngine):
+
+    def __init__(
+        self,
+        pos_embedder: Optional[nn.Module],
+        feat_dim: int,
+        max_points_in_param: int,
+        num_queries: int,
+    ):
+        super().__init__(None, feat_dim, max_points_in_param, num_queries)
+        self.query_embed = nn.Embedding(self.num_queries, feat_dim)
+
+    def forward(self, data: Data) -> tuple[Tensor]:
+        return (
+            torch.zeros(
+                data.batch_size,
+                self.num_queries,
+                self.max_points_in_param,
+                3,
+                device=data.pos.device,
+                requires_grad=False,
+            ),
+            self.query_embed.weight.unsqueeze(0)
+            .expand(data.batch_size, -1, -1)
+            .permute(0, 2, 1),
+        )
+
+
+def build_query_engine(
+    query_type: str,
+    pos_embedder: Optional[nn.Module],
+    feat_dim: int,
+    max_points_in_param: int,
+    num_queries: int,
+) -> QueryEngine:
+    if query_type == "point_fps":
+        return PointFPSQueryEngine(
+            pos_embedder, feat_dim, max_points_in_param, num_queries
+        )
+    elif query_type == "learned":
+        return LearnedQueryEngine(None, feat_dim, max_points_in_param, num_queries)
+    else:
+        raise ValueError(f"Unknown query type {query_type}")

pi3detr/models/transformer.py

@@ -0,0 +1,399 @@
+"""
+Adapted code from Meta's DETR.
+"""
+
+import copy
+from typing import Optional, List
+import torch
+import torch.nn.functional as F
+from torch import nn, Tensor
+
+
+class Transformer(nn.Module):
+
+    def __init__(
+        self,
+        enc_dim: int = 256,
+        dec_dim: int = 256,
+        nhead: int = 8,
+        num_encoder_layers: int = 6,
+        num_decoder_layers: int = 6,
+        enc_dim_feedforward: int = 2048,
+        dec_dim_feedforward: int = 2048,
+        enc_dropout: float = 0.1,
+        dec_dropout: float = 0.1,
+        activation: str = "relu",
+        normalize_before: bool = False,
+        return_intermediate_dec: bool = False,
+    ):
+        super().__init__()
+        encoder_layer = TransformerEncoderLayer(
+            enc_dim,
+            nhead,
+            enc_dim_feedforward,
+            enc_dropout,
+            activation,
+            normalize_before,
+        )
+        encoder_norm = nn.LayerNorm(enc_dim) if normalize_before else None
+        self.encoder = TransformerEncoder(
+            encoder_layer, num_encoder_layers, encoder_norm
+        )
+        if enc_dim != dec_dim:
+            self.enc_to_dec_proj = nn.Linear(enc_dim, dec_dim)
+        else:
+            self.enc_to_dec_proj = nn.Identity()
+        decoder_layer = TransformerDecoderLayer(
+            dec_dim,
+            nhead,
+            dec_dim_feedforward,
+            dec_dropout,
+            activation,
+            normalize_before,
+        )
+        decoder_norm = nn.LayerNorm(dec_dim)
+        self.decoder = TransformerDecoder(
+            decoder_layer,
+            num_decoder_layers,
+            decoder_norm,
+            return_intermediate=return_intermediate_dec,
+        )
+        self._reset_parameters()
+        self.d_model = dec_dim
+        self.nhead = nhead
+
+    def _reset_parameters(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def forward(
+        self,
+        src: Tensor,
+        mask: Optional[Tensor],
+        query_embed: Tensor,
+        pos_embed: Tensor = None,
+    ) -> Tensor:
+        bs, _, _ = src.shape
+
+        src = src.permute(1, 0, 2)  # (bs, seq, feat) -> (seq, bs, feat)
+        if pos_embed is not None:
+            pos_embed = pos_embed.permute(2, 0, 1)
+
+        memory = self.encoder(src, mask=None, src_key_padding_mask=mask, pos=None)
+        memory = self.enc_to_dec_proj(memory)
+
+        query_embed = query_embed.permute(2, 0, 1)
+        tgt = torch.zeros_like(query_embed)
+        hs = self.decoder(
+            tgt,
+            memory,
+            tgt_mask=None,
+            memory_mask=None,
+            tgt_key_padding_mask=None,
+            memory_key_padding_mask=mask,
+            pos=pos_embed,
+            query_pos=query_embed,
+        )
+        return hs.transpose(1, 2)
+
+
+class TransformerEncoder(nn.Module):
+
+    def __init__(self, encoder_layer, num_layers, norm=None):
+        super().__init__()
+        self.layers = _get_clones(encoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+
+    def forward(
+        self,
+        src,
+        mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+    ):
+        output = src
+
+        for layer in self.layers:
+            output = layer(
+                output,
+                src_mask=mask,
+                src_key_padding_mask=src_key_padding_mask,
+                pos=pos,
+            )
+
+        if self.norm is not None:
+            output = self.norm(output)
+
+        return output
+
+
+class TransformerEncoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation="relu",
+        normalize_before=False,
+    ):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(
+        self,
+        src,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+    ):
+        q = k = self.with_pos_embed(src, pos)
+        src2 = self.self_attn(
+            q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask
+        )[0]
+        src = src + self.dropout1(src2)
+        src = self.norm1(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
+        src = src + self.dropout2(src2)
+        src = self.norm2(src)
+        return src
+
+    def forward_pre(
+        self,
+        src,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+    ):
+        src2 = self.norm1(src)
+        q = k = self.with_pos_embed(src2, pos)
+        src2 = self.self_attn(
+            q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask
+        )[0]
+        src = src + self.dropout1(src2)
+        src2 = self.norm2(src)
+        src2 = self.linear2(self.dropout(self.activation(self.linear1(src2))))
+        src = src + self.dropout2(src2)
+        return src
+
+    def forward(
+        self,
+        src,
+        src_mask: Optional[Tensor] = None,
+        src_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+    ):
+        if self.normalize_before:
+            return self.forward_pre(src, src_mask, src_key_padding_mask, pos)
+        return self.forward_post(src, src_mask, src_key_padding_mask, pos)
+
+
+class TransformerDecoder(nn.Module):
+
+    def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
+        super().__init__()
+        self.layers = _get_clones(decoder_layer, num_layers)
+        self.num_layers = num_layers
+        self.norm = norm
+        self.return_intermediate = return_intermediate
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
+        output = tgt
+
+        intermediate = []
+
+        for layer in self.layers:
+            output = layer(
+                output,
+                memory,
+                tgt_mask=tgt_mask,
+                memory_mask=memory_mask,
+                tgt_key_padding_mask=tgt_key_padding_mask,
+                memory_key_padding_mask=memory_key_padding_mask,
+                pos=pos,
+                query_pos=query_pos,
+            )
+            if self.return_intermediate:
+                intermediate.append(self.norm(output))
+
+        if self.norm is not None:
+            output = self.norm(output)
+            if self.return_intermediate:
+                intermediate.pop()
+                intermediate.append(output)
+
+        if self.return_intermediate:
+            return torch.stack(intermediate)
+
+        return output.unsqueeze(0)
+
+
+class TransformerDecoderLayer(nn.Module):
+
+    def __init__(
+        self,
+        d_model,
+        nhead,
+        dim_feedforward=2048,
+        dropout=0.1,
+        activation="relu",
+        normalize_before=False,
+    ):
+        super().__init__()
+        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
+        # Implementation of Feedforward model
+        self.linear1 = nn.Linear(d_model, dim_feedforward)
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim_feedforward, d_model)
+
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.activation = _get_activation_fn(activation)
+        self.normalize_before = normalize_before
+
+    def with_pos_embed(self, tensor, pos: Optional[Tensor]):
+        return tensor if pos is None else tensor + pos
+
+    def forward_post(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
+        q = k = self.with_pos_embed(tgt, query_pos)
+        tgt2 = self.self_attn(
+            q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt = self.norm1(tgt)
+        tgt2 = self.multihead_attn(
+            query=self.with_pos_embed(tgt, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt = self.norm2(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
+        tgt = tgt + self.dropout3(tgt2)
+        tgt = self.norm3(tgt)
+        return tgt
+
+    def forward_pre(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
+        tgt2 = self.norm1(tgt)
+        q = k = self.with_pos_embed(tgt2, query_pos)
+        tgt2 = self.self_attn(
+            q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask
+        )[0]
+        tgt = tgt + self.dropout1(tgt2)
+        tgt2 = self.norm2(tgt)
+        tgt2 = self.multihead_attn(
+            query=self.with_pos_embed(tgt2, query_pos),
+            key=self.with_pos_embed(memory, pos),
+            value=memory,
+            attn_mask=memory_mask,
+            key_padding_mask=memory_key_padding_mask,
+        )[0]
+        tgt = tgt + self.dropout2(tgt2)
+        tgt2 = self.norm3(tgt)
+        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
+        tgt = tgt + self.dropout3(tgt2)
+        return tgt
+
+    def forward(
+        self,
+        tgt,
+        memory,
+        tgt_mask: Optional[Tensor] = None,
+        memory_mask: Optional[Tensor] = None,
+        tgt_key_padding_mask: Optional[Tensor] = None,
+        memory_key_padding_mask: Optional[Tensor] = None,
+        pos: Optional[Tensor] = None,
+        query_pos: Optional[Tensor] = None,
+    ):
+        if self.normalize_before:
+            return self.forward_pre(
+                tgt,
+                memory,
+                tgt_mask,
+                memory_mask,
+                tgt_key_padding_mask,
+                memory_key_padding_mask,
+                pos,
+                query_pos,
+            )
+        return self.forward_post(
+            tgt,
+            memory,
+            tgt_mask,
+            memory_mask,
+            tgt_key_padding_mask,
+            memory_key_padding_mask,
+            pos,
+            query_pos,
+        )
+
+
+def _get_clones(module, N):
+    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])
+
+
+def _get_activation_fn(activation):
+    """Return an activation function given a string"""
+    if activation == "relu":
+        return F.relu
+    if activation == "gelu":
+        return F.gelu
+    if activation == "glu":
+        return F.glu
+    raise RuntimeError(f"activation should be relu/gelu, not {activation}.")

pi3detr/utils/__init__.py

@@ -0,0 +1,3 @@
+from .config_reader import load_args
+from .layer_utils import no_grad, load_weights
+from .curve_fitter import torch_bezier_curve

pi3detr/utils/config_reader.py

@@ -0,0 +1,21 @@
+from typing import Optional
+from argparse import Namespace
+from types import SimpleNamespace
+import yaml
+
+
+def load_yaml(file_path: str) -> yaml.YAMLObject:
+    with open(file_path, "r") as file:
+        try:
+            return yaml.safe_load(file)
+        except yaml.YAMLError as exc:
+            print(exc)
+            yaml.YAMLError("error reading yaml file")
+
+
+def load_args(config: str, parsed_args: Optional[Namespace] = None) -> SimpleNamespace:
+    args = load_yaml(config)
+    parsed_args = vars(parsed_args) if parsed_args else {}
+    args = args | parsed_args
+    args = SimpleNamespace(**args)
+    return args

pi3detr/utils/curve_fitter.py

@@ -0,0 +1,371 @@
+import numpy as np
+import torch
+import math
+
+
+def torch_arc_points(start, mid, end, num_points=100):
+    """
+    Sample points along a circular arc defined by 3 points in 3D, batched.
+    Inputs:
+        start, mid, end: tensors of shape [B, 3]
+        num_points: number of points sampled along the arc
+    Returns:
+        arc_points: tensor of shape [B, num_points, 3]
+    """
+    B = start.shape[0]
+
+    # 1. Compute circle center and normal vector for each batch
+    v1 = mid - start  # [B,3]
+    v2 = end - start  # [B,3]
+
+    normal = torch.cross(v1, v2, dim=1)  # [B,3]
+    normal_norm = normal.norm(dim=1, keepdim=True).clamp(min=1e-8)
+    normal = normal / normal_norm  # normalize
+
+    mid1 = (start + mid) / 2  # [B,3]
+    mid2 = (start + end) / 2  # [B,3]
+
+    # perpendicular directions in the plane
+    perp1 = torch.cross(normal, v1, dim=1)  # [B,3]
+    perp2 = torch.cross(normal, v2, dim=1)  # [B,3]
+
+    # Solve line intersection for each batch:
+    # Line 1: point mid1, direction perp1
+    # Line 2: point mid2, direction perp2
+    # Solve for t in mid1 + t * perp1 = mid2 + s * perp2
+
+    # Construct matrix A and vector b for least squares
+    A = torch.stack([perp1, -perp2], dim=2)  # [B,3,2]
+    b = (mid2 - mid1).unsqueeze(2)  # [B,3,1]
+
+    # Use torch.linalg.lstsq if available, fallback to pinv:
+    try:
+        t_s = torch.linalg.lstsq(A, b).solution  # [B,2,1]
+    except:
+        # fallback
+        At = A.transpose(1, 2)  # [B,2,3]
+        pinv = torch.linalg.pinv(A)  # [B,2,3]
+        t_s = torch.bmm(pinv, b)  # [B,2,1]
+
+    t = t_s[:, 0, 0]  # [B]
+
+    center = mid1 + (perp1 * t.unsqueeze(1))  # [B,3]
+
+    radius = (start - center).norm(dim=1, keepdim=True)  # [B,1]
+
+    # 2. Define local basis in the arc plane
+    x_axis = (start - center) / radius  # [B,3]
+    y_axis = torch.cross(normal, x_axis, dim=1)  # [B,3]
+
+    # 3. Compute angles function
+    def angle_from_vector(v):
+        x = (v * x_axis).sum(dim=1)  # [B]
+        y = (v * y_axis).sum(dim=1)  # [B]
+        angles = torch.atan2(y, x)  # [-pi, pi]
+        angles = angles % (2 * math.pi)
+        return angles
+
+    theta_start = torch.zeros(B, device=start.device)  # [B], 0 since x_axis is ref
+    theta_end = angle_from_vector(end - center)  # [B]
+    theta_mid = angle_from_vector(mid - center)  # [B]
+
+    # 4. Ensure arc goes the correct way (shortest arc through mid)
+    # Helper function vectorized:
+    def between(a, b, c):
+        # returns bool tensor if b is between a and c going CCW mod 2pi
+        return ((a < b) & (b < c)) | ((c < a) & ((a < b) | (b < c)))
+
+    cond = between(theta_start, theta_mid, theta_end)
+
+    # If not cond, swap start/end angles by adding 2pi to one side
+    # We'll add 2pi to whichever angle is smaller to preserve direction
+    theta_start_new = torch.where(
+        cond,
+        theta_start,
+        torch.where(theta_start < theta_end, theta_start, theta_start + 2 * math.pi),
+    )
+    theta_end_new = torch.where(
+        cond,
+        theta_end,
+        torch.where(theta_end < theta_start, theta_end + 2 * math.pi, theta_end),
+    )
+
+    # 5. Sample angles
+    t_lin = (
+        torch.linspace(0, 1, steps=num_points, device=start.device)
+        .unsqueeze(0)
+        .repeat(B, 1)
+    )  # [B, num_points]
+
+    angles = theta_start_new.unsqueeze(1) + t_lin * (
+        theta_end_new - theta_start_new
+    ).unsqueeze(
+        1
+    )  # [B, num_points]
+    angles = angles % (2 * math.pi)
+
+    # 6. Map back to 3D
+    cos_a = torch.cos(angles).unsqueeze(2)  # [B, num_points, 1]
+    sin_a = torch.sin(angles).unsqueeze(2)  # [B, num_points, 1]
+
+    points = center.unsqueeze(1) + radius.unsqueeze(1) * (
+        cos_a * x_axis.unsqueeze(1) + sin_a * y_axis.unsqueeze(1)
+    )  # [B, num_points, 3]
+
+    return points
+
+
+def torch_circle_fitter(
+    points: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Fits a circle to an arbitrary number of 3D points using least squares.
+
+    Args:
+        points: Tensor of shape (B, N, 3), where B = batch size, N = number of points per batch.
+
+    Returns:
+        center_3d: (B, 3) tensor of circle centers in 3D
+        normal: (B, 3) tensor of normal vectors to the circle's plane
+        radius: (B,) tensor of circle radii
+    """
+    B, N, _ = points.shape
+    mean = points.mean(dim=1, keepdim=True)
+    centered = points - mean
+
+    # PCA via SVD
+    U, S, Vh = torch.linalg.svd(centered)
+    normal = Vh[
+        :, -1, :
+    ]  # last singular vector corresponds to the smallest variance (plane normal)
+
+    # Project to plane
+    x_axis = Vh[:, 0, :]
+    y_axis = Vh[:, 1, :]
+    X = torch.einsum("bij,bj->bi", centered, x_axis)  # (B, N)
+    Y = torch.einsum("bij,bj->bi", centered, y_axis)  # (B, N)
+
+    # Fit circle in 2D: (x - xc)^2 + (y - yc)^2 = r^2
+    A = torch.stack([2 * X, 2 * Y, torch.ones_like(X)], dim=-1)  # (B, N, 3)
+    b = (X**2 + Y**2).unsqueeze(-1)  # (B, N, 1)
+
+    # Solve the least squares system: A @ [xc, yc, c] = b
+    AtA = A.transpose(1, 2) @ A
+    Atb = A.transpose(1, 2) @ b
+    sol = torch.linalg.solve(AtA, Atb).squeeze(-1)  # (B, 3)
+
+    xc, yc, c = sol[:, 0], sol[:, 1], sol[:, 2]
+    radius = torch.sqrt(xc**2 + yc**2 + c)
+
+    # Reconstruct center in 3D
+    center_3d = mean.squeeze(1) + xc.unsqueeze(1) * x_axis + yc.unsqueeze(1) * y_axis
+
+    return center_3d, normal, radius
+
+
+def generate_points_on_circle(center, normal, radius, num_points=100):
+    # Normalize the normal vector
+    normal = normal / np.linalg.norm(normal)
+
+    # Find two orthogonal vectors in the plane of the circle
+    if np.allclose(normal, [0, 0, 1]):
+        u = np.array([1, 0, 0])
+    else:
+        u = np.cross(normal, [0, 0, 1])
+        u = u / np.linalg.norm(u)
+    v = np.cross(normal, u)
+
+    # Generate points on the circle in the plane
+    theta = np.linspace(0, 2 * np.pi, num_points)
+    circle_points = (
+        center
+        + radius * np.outer(np.cos(theta), u)
+        + radius * np.outer(np.sin(theta), v)
+    )
+
+    return circle_points
+
+
+def generate_points_on_circle_torch(
+    center, normal, radius, num_points=100
+) -> torch.Tensor:
+    """
+    Generate points on a circle in 3D space using PyTorch, supporting batching.
+
+    Args:
+        center: Tensor of shape (B, 3), circle centers.
+        normal: Tensor of shape (B, 3), normal vectors to the circle's plane.
+        radius: Tensor of shape (B,), radii of the circles.
+        num_points: Number of points to generate per circle.
+
+    Returns:
+        Tensor of shape (B, num_points, 3), points on the circles.
+    """
+    B = center.shape[0]
+    normal = normal / torch.norm(normal, dim=1, keepdim=True)  # Normalize normals
+
+    # Find two orthogonal vectors in the plane of the circle
+    u = torch.linalg.cross(
+        normal,
+        torch.tensor([0, 0, 1], dtype=normal.dtype, device=normal.device).expand_as(
+            normal
+        ),
+    )
+    u = torch.where(
+        torch.norm(u, dim=1, keepdim=True) > 1e-6,
+        u,
+        torch.tensor([1, 0, 0], dtype=normal.dtype, device=normal.device).expand_as(
+            normal
+        ),
+    )
+    u = u / torch.norm(u, dim=1, keepdim=True)
+    v = torch.linalg.cross(normal, u)
+
+    # Generate points on the circle in the plane
+    theta = (
+        torch.linspace(0, 2 * torch.pi, num_points, device=center.device)
+        .unsqueeze(0)
+        .repeat(B, 1)
+    )
+    circle_points = (
+        center.unsqueeze(1)
+        + radius.unsqueeze(1).unsqueeze(2)
+        * torch.cos(theta).unsqueeze(2)
+        * u.unsqueeze(1)
+        + radius.unsqueeze(1).unsqueeze(2)
+        * torch.sin(theta).unsqueeze(2)
+        * v.unsqueeze(1)
+    )
+
+    return circle_points
+
+
+def torch_bezier_curve(
+    control_points: torch.Tensor, num_points: int = 100
+) -> torch.Tensor:
+    control_points = control_points.float()
+    t = (torch.linspace(0, 1, num_points).unsqueeze(-1).unsqueeze(-1)).to(
+        control_points.device
+    )  # shape [1, num_points, 1]
+    B = (
+        (1 - t) ** 3 * control_points[:, 0]
+        + 3 * (1 - t) ** 2 * t * control_points[:, 1]
+        + 3 * (1 - t) * t**2 * control_points[:, 2]
+        + t**3 * control_points[:, 3]
+    )
+    # Transpose the first two dimensions to get the shape (batch_size, num_points, 3)
+    B = B.transpose(0, 1)
+
+    return B
+
+
+def torch_line_points(
+    start_points: torch.Tensor, end_points: torch.Tensor, num_points: int = 100
+) -> torch.Tensor:
+    weights = (
+        torch.linspace(0, 1, num_points)
+        .unsqueeze(0)
+        .unsqueeze(-1)
+        .to(start_points.device)
+    )
+    line_points = (1 - weights) * start_points.unsqueeze(
+        1
+    ) + weights * end_points.unsqueeze(1)
+    return line_points
+
+
+def fit_line(points: torch.Tensor, K: int = 100) -> torch.Tensor:
+    """
+    Fit a line to 3D points and sample K points along it.
+    """
+    assert points.ndim == 2 and points.shape[1] == 3, "Input must be [N, 3]"
+
+    # Step 1: Center the points
+    mean = points.mean(dim=0, keepdim=True)
+    centered = points - mean
+
+    # Step 2: SVD
+    U, S, Vh = torch.linalg.svd(centered, full_matrices=False)
+    direction = Vh[0]  # First principal component
+
+    # Step 3: Project points onto the line to get min/max
+    projections = torch.matmul(centered, direction)
+    t_min, t_max = projections.min(), projections.max()
+
+    # Step 4: Sample along the line
+    t_vals = torch.linspace(t_min, t_max, K).to(points.device)
+    fitted_points = mean + t_vals[:, None] * direction
+
+    return fitted_points
+
+
+def fit_cubic_bezier(points_3d: torch.Tensor) -> torch.Tensor:
+    """
+    Fit a cubic Bézier curve to 3D points while fixing the start and end points.
+
+    Args:
+        points_3d: (N, 3) Tensor of 3D arc points.
+
+    Returns:
+        bezier_pts: Tensor of 4 control points (P0, P1, P2, P3), shape (4, 3)
+    """
+    if not isinstance(points_3d, torch.Tensor):
+        points_3d = torch.tensor(points_3d, dtype=torch.float32)
+
+    n = len(points_3d)
+
+    if n < 4:
+        raise ValueError("At least 4 points are required to fit a cubic Bézier curve.")
+
+    device = points_3d.device
+
+    # Fixed start and end points
+    P0 = points_3d[0]
+    P3 = points_3d[-1]
+
+    # Normalize parameter t
+    t = torch.linspace(0, 1, n, device=device)
+
+    # Bernstein basis functions for cubic Bézier
+    def bernstein(t):
+        b0 = (1 - t) ** 3
+        b1 = 3 * (1 - t) ** 2 * t
+        b2 = 3 * (1 - t) * t**2
+        b3 = t**3
+        return torch.stack([b0, b1, b2, b3], dim=1)  # (n, 4)
+
+    B = bernstein(t)
+
+    # Initial guess for P1 and P2 (based on tangents)
+    P1_init = P0 + (points_3d[1] - P0) * 1.5
+    P2_init = P3 + (points_3d[-2] - P3) * 1.5
+
+    # Optimization parameters - make them require gradients
+    P1 = P1_init.clone().detach().requires_grad_(True)
+    P2 = P2_init.clone().detach().requires_grad_(True)
+
+    # Optimizer
+    optimizer = torch.optim.LBFGS([P1, P2], max_iter=100, line_search_fn="strong_wolfe")
+
+    def closure():
+        optimizer.zero_grad()
+
+        # Compute Bézier curve
+        curve = (
+            B[:, 0].unsqueeze(1) * P0
+            + B[:, 1].unsqueeze(1) * P1
+            + B[:, 2].unsqueeze(1) * P2
+            + B[:, 3].unsqueeze(1) * P3
+        )
+
+        # Compute loss (mean squared error)
+        loss = torch.mean((curve - points_3d) ** 2)
+        loss.backward()
+        return loss
+
+    # Optimize
+    optimizer.step(closure)
+
+    # Return control points
+    with torch.no_grad():
+        return torch.stack([P0, P1, P2, P3])

pi3detr/utils/layer_utils.py

@@ -0,0 +1,22 @@
+from typing import Optional
+import torch
+from torch import nn
+
+
+def load_weights(
+    model: nn.Module, ckpt_path: str, dont_load: Optional[list[str]] = []
+) -> None:
+    ckpt = torch.load(ckpt_path, weights_only=False)
+    state_dict = {}
+    for k, v in ckpt["state_dict"].items():
+        if not any([dl in k for dl in dont_load]):
+            state_dict[k] = v
+        else:
+            print(f"Didn't load {k}")
+    model.load_state_dict(state_dict, strict=False)
+    print(f"Loaded checkpoint: {ckpt_path}")
+
+
+def no_grad(model: nn.Module) -> None:
+    for param in model.parameters():
+        param.requires_grad = False

pi3detr/utils/postprocessing.py

@@ -0,0 +1,543 @@
+import numpy as np
+import torch
+from torch import Tensor
+from torch_geometric.data import Data
+from scipy.spatial import cKDTree  # faster than KDTree
+
+from .curve_fitter import (
+    fit_cubic_bezier,
+    torch_bezier_curve,
+    torch_circle_fitter,
+    generate_points_on_circle_torch,
+    torch_arc_points,
+    fit_line,
+)
+
+
+def snap_and_fit_curves(
+    data: Data,
+) -> Data:
+    """
+    Snap polylines to nearest point cloud points and fit geometric curves based on predicted classes.
+
+    This function performs two main operations:
+    1. Snaps each polyline vertex to its nearest neighbor in the point cloud
+    2. Fits the appropriate geometric curve (line, circle, arc, or B-spline) based on the predicted class
+
+    Class mapping:
+        0: Background (no processing, kept as-is)
+        1: B-spline (cubic Bezier curve fitting)
+        2: Line (linear regression fitting)
+        3: Circle (3D circle fitting)
+        4: Arc (circular arc through 3 points)
+
+    Args:
+        data (Data): PyTorch Geometric Data object containing:
+            - pos (Tensor): Point cloud coordinates [P, 3]
+            - polylines (Tensor): Raw polyline predictions [M, K, 3]
+            - polyline_class (Tensor): Predicted classes for each polyline [M]
+
+    Returns:
+        Data: Cloned Data object with fitted polylines replacing the original polylines.
+              All other attributes remain unchanged.
+
+    Note:
+        - Robust error handling: falls back to original polyline if fitting fails
+        - Validates output shapes and numerical stability (NaN/Inf checks)
+        - Requires minimum points per curve type (e.g., 3 for circles, 4 for B-splines)
+    """
+    point_cloud = data.pos
+    polylines = data.polylines
+    polyline_classes = data.polyline_class
+    M, K, _ = polylines.shape
+    snapped_and_fitted = torch.zeros_like(polylines)
+
+    for i, cls in enumerate(polyline_classes):
+        if cls == 0:
+            snapped_and_fitted[i] = polylines[i]  # Keep original for class 0
+            continue
+
+        try:
+            # Snap the polyline to the nearest point in the point cloud
+            distances = torch.cdist(polylines[i], point_cloud)
+            nearest_idx = distances.argmin(dim=1)
+            nn_points = point_cloud[nearest_idx]
+
+            # Safety check: ensure we have valid points
+            if (
+                len(nn_points) == 0
+                or torch.any(torch.isnan(nn_points))
+                or torch.any(torch.isinf(nn_points))
+            ):
+                snapped_and_fitted[i] = polylines[i]
+                continue
+
+            new_curve = None
+
+            if cls == 1:  # BSpline
+                try:
+                    if len(nn_points) < 4:
+                        # Not enough points for cubic Bezier, fallback to original
+                        new_curve = polylines[i]
+                    else:
+                        bezier_pts = fit_cubic_bezier(nn_points)
+                        new_curve = torch_bezier_curve(
+                            bezier_pts.unsqueeze(0), K
+                        ).squeeze(0)
+
+                        # Validate output shape and values
+                        if (
+                            new_curve.shape != (K, 3)
+                            or torch.any(torch.isnan(new_curve))
+                            or torch.any(torch.isinf(new_curve))
+                        ):
+                            new_curve = polylines[i]
+
+                except Exception:
+                    new_curve = polylines[i]
+
+            elif cls == 2:  # Line
+                try:
+                    if len(nn_points) < 2:
+                        new_curve = polylines[i]
+                    else:
+                        new_curve = fit_line(nn_points, K)
+
+                        # Validate output shape and values
+                        if (
+                            new_curve.shape != (K, 3)
+                            or torch.any(torch.isnan(new_curve))
+                            or torch.any(torch.isinf(new_curve))
+                        ):
+                            new_curve = polylines[i]
+                except Exception:
+                    new_curve = polylines[i]
+
+            elif cls == 3:  # Circle
+                try:
+                    # Check if we have enough unique points for circle fitting
+                    unique_points = torch.unique(nn_points, dim=0)
+                    if len(unique_points) < 3:
+                        new_curve = polylines[i]
+                    else:
+                        center, normal, radius = torch_circle_fitter(
+                            nn_points.unsqueeze(0)
+                        )
+
+                        # Validate circle parameters
+                        if (
+                            torch.any(torch.isnan(center))
+                            or torch.any(torch.isnan(normal))
+                            or torch.any(torch.isnan(radius))
+                            or torch.any(torch.isinf(center))
+                            or torch.any(torch.isinf(normal))
+                            or torch.any(torch.isinf(radius))
+                            or radius <= 0
+                        ):
+                            new_curve = polylines[i]
+                        else:
+                            new_curve = generate_points_on_circle_torch(
+                                center, normal, radius, K
+                            ).squeeze(0)
+
+                            # Validate output shape and values
+                            if (
+                                new_curve.shape != (K, 3)
+                                or torch.any(torch.isnan(new_curve))
+                                or torch.any(torch.isinf(new_curve))
+                            ):
+                                new_curve = polylines[i]
+                except Exception:
+                    new_curve = polylines[i]
+
+            elif cls == 4:  # Arc
+                try:
+                    if len(nn_points) < 3:
+                        new_curve = polylines[i]
+                    else:
+                        start_pt = nn_points[0].unsqueeze(0)
+                        mid_pt = nn_points[len(nn_points) // 2].unsqueeze(0)
+                        end_pt = nn_points[-1].unsqueeze(0)
+
+                        new_curve = torch_arc_points(
+                            start_pt, mid_pt, end_pt, K
+                        ).squeeze(0)
+
+                        # Validate output shape and values
+                        if (
+                            new_curve.shape != (K, 3)
+                            or torch.any(torch.isnan(new_curve))
+                            or torch.any(torch.isinf(new_curve))
+                        ):
+                            new_curve = polylines[i]
+                except Exception:
+                    new_curve = polylines[i]
+
+            else:
+                # Unknown class, keep original
+                new_curve = polylines[i]
+
+            # Final safety check
+            if new_curve is not None and new_curve.shape == (K, 3):
+                snapped_and_fitted[i] = new_curve
+            else:
+                snapped_and_fitted[i] = polylines[i]
+
+        except Exception:
+            # If anything goes wrong, fallback to original polyline
+            snapped_and_fitted[i] = polylines[i]
+
+    output = data.clone()
+    output.polylines = snapped_and_fitted
+    return output
+
+
+def filter_predictions(pred_data: Data, thresholds: list[float]) -> Data:
+    """
+    Filter predictions based on class-specific confidence thresholds.
+
+    Removes polylines whose confidence scores fall below the specified threshold
+    for their predicted class. This is typically used as a post-processing step
+    to remove low-confidence predictions before further analysis.
+
+    Args:
+        pred_data (Data): PyTorch Geometric Data object containing:
+            - pos (Tensor): Point cloud coordinates [P, 3]
+            - polyline_class (Tensor): Predicted classes [N]
+            - polyline_score (Tensor): Confidence scores [N]
+            - polylines (Tensor): Polyline coordinates [N, K, 3]
+            - query_xyz (Tensor, optional): Query coordinates [N, 3]
+        thresholds (list[float]): Confidence thresholds for each class.
+                                 Length must match the number of classes.
+                                 thresholds[i] is the minimum confidence for class i.
+
+    Returns:
+        Data: Filtered Data object containing only polylines that meet their
+              class-specific confidence thresholds. Maintains the same structure
+              as input but with potentially fewer polylines.
+
+    Example:
+        # Keep only polylines with confidence > 0.5 for class 0, > 0.7 for class 1, etc.
+        filtered = filter_predictions(data, [0.5, 0.7, 0.6, 0.8])
+    """
+    mask = (
+        pred_data.polyline_score
+        >= torch.tensor(thresholds, device=pred_data.pos.device)[
+            pred_data.polyline_class
+        ]
+    )
+    filtered_data = Data(
+        pos=pred_data.pos,
+        polyline_class=pred_data.polyline_class[mask],
+        polyline_score=pred_data.polyline_score[mask],
+        polylines=pred_data.polylines[mask],
+        query_xyz=(
+            pred_data.query_xyz[mask] if hasattr(pred_data, "query_xyz") else None
+        ),
+    )
+
+    return filtered_data
+
+
+def iou_filter_point_based(
+    pred_data,
+    iou_threshold: float = 0.6,
+    background_class: int = 0,
+):
+    """
+    Efficient per-class Non-Maximum Suppression using IoU computed on point cloud indices.
+
+    This optimized NMS implementation:
+    1. Snaps all polyline vertices to nearest point cloud neighbors
+    2. Computes IoU based on overlapping point cloud indices (not 3D distances)
+    3. Applies greedy NMS within each class, keeping highest-scoring polylines
+    4. Uses optimized data structures (cKDTree, sorted arrays) for speed
+
+    Algorithm details:
+    - Single batched nearest neighbor query for all valid vertices
+    - IoU = |intersection| / |union| of snapped point indices
+    - Polylines ordered by: score (desc) → #snapped_points (desc) → index (asc)
+    - Background class polylines are never removed
+    - Polylines with no valid snapped points are dropped
+
+    Args:
+        pred_data (Data): PyTorch Geometric Data object with polyline predictions
+        iou_threshold (float, optional): IoU threshold for suppression. Default: 0.6
+        background_class (int, optional): Class ID to exclude from NMS. Default: 0
+
+    Returns:
+        Data: Filtered Data object with overlapping polylines removed per class.
+              Maintains same structure with potentially fewer polylines.
+
+    Performance:
+        Significantly faster than distance-based methods due to:
+        - Batched spatial queries (cKDTree)
+        - Integer set operations (np.intersect1d)
+        - Minimal Python loops
+    """
+    data = pred_data.clone()
+
+    polylines: torch.Tensor = data.polylines  # (N, M, 3)
+    classes: torch.Tensor = data.polyline_class  # (N,)
+    pc: torch.Tensor = data.pos  # (P, 3)
+    scores = getattr(data, "polyline_score", None)
+
+    device = polylines.device
+    N = polylines.shape[0]
+    if N == 0 or pc.shape[0] == 0:
+        return data
+
+    # ---- helpers ----
+    def valid_mask(pts_t: torch.Tensor) -> torch.Tensor:
+        finite = torch.isfinite(pts_t).all(dim=-1)
+        non_zero = pts_t.abs().sum(dim=-1) > 0
+        return finite & non_zero
+
+    # ---- gather all valid vertices once (batched) ----
+    # We'll collect (poly_idx, vertex_xyz) over all non-background curves.
+    poly_indices_list = []
+    all_vertices = []
+
+    bg = int(background_class)
+    for i in range(N):
+        if int(classes[i].item()) == bg:
+            continue
+        vm = valid_mask(polylines[i])
+        if vm.any():
+            pts = polylines[i][vm].detach().cpu().numpy()
+            if pts.size > 0:
+                all_vertices.append(pts)
+                poly_indices_list.append(np.full((pts.shape[0],), i, dtype=np.int32))
+
+    if len(all_vertices) == 0:
+        # nothing to snap; everything gets dropped
+        keep_mask = torch.zeros(N, dtype=torch.bool, device=device)
+        data.polylines = data.polylines[keep_mask]
+        data.polyline_class = data.polyline_class[keep_mask]
+        if hasattr(data, "polyline_score") and data.polyline_score is not None:
+            data.polyline_score = data.polyline_score[keep_mask]
+        if hasattr(data, "query_xyz") and data.query_xyz is not None:
+            data.query_xyz = data.query_xyz[keep_mask]
+        return data
+
+    all_vertices = np.concatenate(all_vertices, axis=0)  # (T, 3)
+    owner_poly = np.concatenate(poly_indices_list, axis=0)  # (T,)
+
+    # ---- one cKDTree query for all vertices ----
+    pc_np = pc.detach().cpu().numpy()
+    tree = cKDTree(pc_np)
+    # Use parallel workers if SciPy supports it (falls back silently otherwise)
+    nn_idx = tree.query(all_vertices, workers=-1)[1].astype(np.int64)  # (T,)
+
+    # ---- split back to per-curve snapped unique index arrays (sorted) ----
+    snapped_arrays = [None] * N
+    set_sizes = torch.zeros(N, dtype=torch.long, device=device)
+
+    # group indices by polyline using numpy argsort
+    order = np.argsort(owner_poly, kind="mergesort")
+    owner_sorted = owner_poly[order]
+    nn_sorted = nn_idx[order]
+
+    # find segment starts for each unique polyline id
+    unique_ids, starts = np.unique(owner_sorted, return_index=True)
+    # append end sentinel
+    starts = np.append(starts, owner_sorted.shape[0])
+
+    for k in range(len(unique_ids)):
+        i = int(unique_ids[k])
+        seg = nn_sorted[starts[k] : starts[k + 1]]
+        if seg.size == 0:
+            snapped_arrays[i] = np.empty((0,), dtype=np.int64)
+            continue
+        uniq = np.unique(seg)  # already sorted
+        snapped_arrays[i] = uniq
+        set_sizes[i] = uniq.size
+
+    # For background curves or curves with no valid vertices, ensure empty arrays
+    for i in range(N):
+        if snapped_arrays[i] is None:
+            snapped_arrays[i] = np.empty((0,), dtype=np.int64)
+
+    # fallback scores: prefer more snapped support
+    if scores is None:
+        scores = set_sizes.to(torch.float)
+
+    keep_mask = torch.ones(N, dtype=torch.bool, device=device)
+
+    # ---- per-class greedy NMS (IoU via fast array intersection) ----
+    target_classes = torch.unique(classes[classes != background_class]).tolist()
+    for cls in target_classes:
+        cls_inds = torch.where(classes == cls)[0].tolist()
+        if not cls_inds:
+            continue
+
+        # order by (score desc, size desc, index asc)
+        cls_order = sorted(
+            cls_inds,
+            key=lambda idx: (
+                -float(scores[idx].item()),
+                -int(set_sizes[idx].item()),
+                idx,
+            ),
+        )
+
+        suppressed = set()
+        for i_idx in cls_order:
+            if i_idx in suppressed:
+                continue
+
+            A = snapped_arrays[i_idx]
+            if A.size == 0:
+                suppressed.add(i_idx)
+                continue
+
+            lenA = A.size
+            for j_idx in cls_order:
+                if j_idx <= i_idx or j_idx in suppressed:
+                    continue
+                B = snapped_arrays[j_idx]
+                if B.size == 0:
+                    suppressed.add(j_idx)
+                    continue
+
+                # fast intersection of two sorted int arrays
+                inter = np.intersect1d(A, B, assume_unique=True).size
+                union = lenA + B.size - inter
+                if union == 0:
+                    continue
+                if (inter / union) > iou_threshold:
+                    suppressed.add(j_idx)
+
+        if suppressed:
+            keep_mask[list(suppressed)] = False
+
+    # ---- filter aligned fields ----
+    data.polylines = data.polylines[keep_mask]
+    data.polyline_class = data.polyline_class[keep_mask]
+    if hasattr(data, "polyline_score") and data.polyline_score is not None:
+        data.polyline_score = data.polyline_score[keep_mask]
+    if hasattr(data, "query_xyz") and data.query_xyz is not None:
+        data.query_xyz = data.query_xyz[keep_mask]
+
+    return data
+
+
+def iou_filter_predictions(
+    data: Data,
+    iou_threshold: float = 0.6,
+    tol: float = 1e-2,
+) -> Data:
+    """
+    Remove overlapping polylines within each class using point-to-point distance IoU.
+
+    Performs class-wise Non-Maximum Suppression to eliminate redundant predictions:
+    1. Filters out invalid points (NaN, Inf, near-zero)
+    2. Computes pairwise point distances between polylines of the same class
+    3. Calculates IoU based on points within distance tolerance
+    4. Removes lower-scoring polylines when IoU exceeds threshold
+    5. Protects "lonely" polylines (minimal overlap) from removal
+
+    IoU Calculation:
+        - overlap_i = number of points in polyline_i within tolerance of polyline_j
+        - overlap_j = number of points in polyline_j within tolerance of polyline_i
+        - intersection = min(overlap_i, overlap_j)
+        - union = len(polyline_i) + len(polyline_j) - intersection
+        - IoU = intersection / union
+
+    Args:
+        data (Data): PyTorch Geometric Data object containing:
+            - polylines (Tensor): Polyline coordinates [N, P, 3]
+            - polyline_class (Tensor): Class predictions [N]
+            - polyline_score (Tensor): Confidence scores [N]
+            - query_xyz (Tensor, optional): Query coordinates [N, 3]
+        iou_threshold (float, optional): IoU threshold for duplicate removal. Default: 0.6
+        tol (float, optional): Distance tolerance for point overlap detection. Default: 1e-2
+
+    Returns:
+        Data: Filtered Data object with overlapping polylines removed.
+              Background class (0) polylines are never removed.
+
+    Note:
+        - Processes polylines in descending score order for deterministic results
+        - Requires significant overlap (≥2 points, ≥10% of smaller polyline) before considering removal
+        - More computationally expensive than index-based methods but handles arbitrary point clouds
+    """
+    polylines = data.polylines
+    polyline_class = data.polyline_class
+    scores = data.polyline_score
+
+    # Precompute valid points for all polylines
+    valid_pts = []
+    for poly in polylines:
+        mask = ~torch.isnan(poly).any(dim=1) & (
+            ~torch.isinf(poly).any(dim=1) & (torch.norm(poly, dim=1) > 1e-6)
+        )
+        valid_pts.append(poly[mask])
+
+    remove_set = set()
+
+    # Process each class independently
+    for cls in torch.unique(polyline_class):
+        if cls == 0:  # Skip background
+            continue
+        # Get indices for this class
+        class_mask = polyline_class == cls
+        class_indices = torch.where(class_mask)[0]
+        if len(class_indices) < 2:
+            continue
+
+        # Sort by score descending, then index ascending for determinism
+        sorted_indices = sorted(
+            class_indices.tolist(), key=lambda idx: (-scores[idx].item(), idx)
+        )
+
+        # Compare pairs in sorted order
+        for i, idx_i in enumerate(sorted_indices):
+            if idx_i in remove_set:
+                continue
+            pts_i = valid_pts[idx_i]
+            if len(pts_i) == 0:
+                continue
+            for j in range(i + 1, len(sorted_indices)):
+                idx_j = sorted_indices[j]
+                if idx_j in remove_set:
+                    continue
+                pts_j = valid_pts[idx_j]
+                if len(pts_j) == 0:
+                    continue
+
+                # Compute point-wise distances
+                dists = torch.cdist(pts_i, pts_j)
+                # Calculate overlaps
+                overlap_i = (dists.min(dim=1).values < tol).sum().item()
+                overlap_j = (dists.min(dim=0).values < tol).sum().item()
+                min_points = min(len(pts_i), len(pts_j))
+                # Skip if not significant overlap
+                if (
+                    overlap_i < 2
+                    or overlap_j < 2
+                    or min(overlap_i, overlap_j) < 0.1 * min_points
+                ):
+                    continue
+
+                # Calculate IoU
+                intersection = min(overlap_i, overlap_j)
+                union = len(pts_i) + len(pts_j) - intersection
+                iou = intersection / union if union > 0 else 0.0
+                if iou > iou_threshold:
+                    # Always remove lower-scoring polyline
+                    remove_set.add(idx_j)
+
+    # Create keep mask (protects lonely lines)
+    keep_mask = torch.ones(len(polylines), dtype=torch.bool)
+    for idx in remove_set:
+        keep_mask[idx] = False
+
+    # Apply filtering
+    data.polylines = polylines[keep_mask]
+    data.polyline_class = polyline_class[keep_mask]
+    data.polyline_score = scores[keep_mask]
+    if hasattr(data, "query_xyz"):
+        data.query_xyz = data.query_xyz[keep_mask]
+
+    return data

pi3detr/utils/viz.py

@@ -0,0 +1,12 @@
+def hex_to_rgb(hex_color: str) -> tuple:
+    """
+    Convert hex color string to RGB tuple.
+
+    Args:
+        hex_color: Hex color string (e.g., "#FF5733")
+
+    Returns:
+        Tuple of RGB values (0-1 range)
+    """
+    hex_color = hex_color.lstrip("#")
+    return tuple(int(hex_color[i : i + 2], 16) / 255.0 for i in (0, 2, 4))

requirements.txt

@@ -0,0 +1,46 @@
+# ==== Repos for CUDA 12.1 PyTorch + PyG wheels ====
+--index-url https://download.pytorch.org/whl/cu121
+--extra-index-url https://pypi.org/simple
+--find-links https://data.pyg.org/whl/torch-2.5.1%2Bcu121.html
+
+# ==== Core: PyTorch CUDA 12.1 ====
+torch==2.5.1+cu121
+
+# ==== PyTorch Geometric stack (built for torch 2.5 + cu121) ====
+pyg-lib==0.4.0+pt25cu121
+torch-scatter==2.1.2+pt25cu121
+torch-sparse==0.6.18+pt25cu121
+torch-cluster==1.6.3+pt25cu121
+torch-spline-conv==1.2.2+pt25cu121
+torch-geometric==2.6.1
+
+# ==== Vision / geometry extras ====
+kornia==0.8.0
+opencv-python==4.11.0.86
+open3d==0.19.0
+polyscope==2.3.0
+trimesh==4.6.8
+timm==1.0.14
+spconv-cu121==2.3.8
+fpsample==0.3.3
+
+# ==== SciPy stack ====
+# Use NumPy >=2.0 for SciPy 1.15.x; leave upper bound open for compatibility with CUDA wheels.
+numpy>=2.0
+scipy==1.15.2
+scikit-learn==1.6.1
+matplotlib==3.10.0
+
+# ==== Jupyter / developer tools ====
+ipython>=8.20
+ipykernel==6.29.5
+ipywidgets==8.1.5
+black==25.1.0
+tqdm>=4.66
+tensorboard==2.19.0
+pytorch-lightning==2.5.0.post0
+
+# ==== Hugging Face Space ====
+gradio==5.49.1
+plotly==6.3.1
+plyfile==1.1.2