point_sam/inference.py

"""High-level inference API for Point-SAM.

This module provides a clean, hydra-free interface for running Point-SAM
segmentation on point clouds loaded from PLY or PCD files.
"""

import os
import warnings
from typing import Union, Tuple, Optional

import numpy as np
import torch
import torch.nn as nn
from safetensors.torch import load_model as load_safetensors_model

from .model.pc_encoder import PatchEmbed, PointCloudEncoder
from .model.pc_sam import PointCloudSAM
from .model.prompt_encoder import MaskEncoder, PointEncoder
from .model.mask_decoder import MaskDecoder
from .model.transformer import TwoWayTransformer
from .utils.torch_utils import replace_with_fused_layernorm


def _load_ply_ascii(filename: str) -> np.ndarray:
    """Load an ASCII PLY file with xyzrgb columns."""
    with open(filename, "r") as rf:
        num_of_points = None
        while True:
            line = rf.readline()
            if not line:
                break
            if "end_header" in line:
                break
            if "element vertex" in line:
                num_of_points = int(line.split()[2])
        if num_of_points is None:
            raise ValueError(f"Could not parse vertex count from PLY header: {filename}")
        points = np.zeros([num_of_points, 6], dtype=np.float32)
        for i in range(num_of_points):
            point = rf.readline().split()
            if len(point) < 6:
                raise ValueError(
                    f"Line {i} in PLY has fewer than 6 values ({len(point)})."
                )
            points[i] = [float(p) for p in point[:6]]
    return points


def _load_pcd_ascii(filename: str) -> np.ndarray:
    """Load an ASCII PCD file with xyzrgb columns."""
    with open(filename, "r") as rf:
        header_ended = False
        num_of_points = None
        data_mode = None
        while True:
            line = rf.readline()
            if not line:
                break
            line = line.strip()
            if line.startswith("POINTS "):
                num_of_points = int(line.split()[1])
            if line.startswith("DATA "):
                data_mode = line.split()[1]
                header_ended = True
                break
        if num_of_points is None:
            raise ValueError(f"Could not parse POINTS from PCD header: {filename}")
        if data_mode != "ascii":
            raise ValueError(f"Only ASCII PCD is supported; got DATA {data_mode}")
        points = np.zeros([num_of_points, 6], dtype=np.float32)
        for i in range(num_of_points):
            point = rf.readline().split()
            if len(point) < 6:
                # Some PCD files store x y z rgb as single float — try unpacking
                if len(point) == 4:
                    # x y z rgb (packed) — unpack rgb into r g b
                    rgb_packed = float(point[3])
                    rgb_int = int(rgb_packed)
                    r = ((rgb_int >> 16) & 0xFF)
                    g = ((rgb_int >> 8) & 0xFF)
                    b = (rgb_int & 0xFF)
                    points[i] = [float(point[0]), float(point[1]), float(point[2]), r, g, b]
                else:
                    raise ValueError(
                        f"Line {i} in PCD has fewer than 6 values ({len(point)})."
                    )
            else:
                points[i] = [float(p) for p in point[:6]]
    return points


def load_pointcloud(
    filepath: str,
    normalize: bool = True,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Load a point cloud from a PLY or PCD file.

    Args:
        filepath: Path to .ply or .pcd file.
        normalize: Whether to normalize coordinates to a unit sphere in [-1, 1].

    Returns:
        coords: [N, 3] numpy array of coordinates.
        rgb: [N, 3] numpy array of colors in [0, 255].
        original_coords: [N, 3] un-normalized coordinates.
    """
    ext = os.path.splitext(filepath)[1].lower()
    if ext == ".ply":
        points = _load_ply_ascii(filepath)
    elif ext == ".pcd":
        points = _load_pcd_ascii(filepath)
    else:
        raise ValueError(f"Unsupported file extension: {ext}. Use .ply or .pcd")

    original_coords = points[:, :3].copy()
    rgb = points[:, 3:6].copy()

    # If colors look very small, they may already be in [0, 1]
    if rgb.max() <= 1.0 + 1e-6:
        rgb = rgb * 255.0

    if normalize:
        coords = original_coords - original_coords.mean(axis=0)
        max_norm = np.linalg.norm(coords, axis=1).max()
        if max_norm > 1e-8:
            coords = coords / max_norm
    else:
        coords = original_coords

    return coords, rgb, original_coords


def build_point_sam(
    variant: str = "large",
    embed_dim: int = 256,
    device: str = "cuda",
    use_fused_layernorm: bool = False,
) -> PointCloudSAM:
    """Build a Point-SAM model from scratch (no hydra/omegaconf required).

    Args:
        variant: Model size — "large" or "giant".
        embed_dim: Embedding dimension for the decoder.
        device: torch device to place the model on.
        use_fused_layernorm: Whether to replace LayerNorm with apex FusedLayerNorm.
            Requires apex to be installed.

    Returns:
        PointCloudSAM model on the requested device.
    """
    import timm

    if variant == "large":
        model_name = "eva02_large_patch14_448"
        num_patches = 1024
        patch_size = 256
        prompt_iters = 5
    elif variant == "giant":
        model_name = "eva_giant_patch14_560"
        num_patches = 512
        patch_size = 64
        prompt_iters = 10
    else:
        raise ValueError(f"Unknown variant: {variant}. Choose 'large' or 'giant'.")

    # Build encoder
    transformer_encoder = timm.create_model(model_name, pretrained=False)
    patch_embed = PatchEmbed(
        in_channels=6,
        out_channels=512,
        num_patches=num_patches,
        patch_size=patch_size,
    )
    pc_encoder = PointCloudEncoder(
        patch_embed=patch_embed,
        transformer=transformer_encoder,
        embed_dim=embed_dim,
    )

    # Build prompt encoder
    mask_encoder = MaskEncoder(embed_dim=embed_dim)

    # Build decoder
    transformer_decoder = TwoWayTransformer(
        depth=2,
        embedding_dim=embed_dim,
        num_heads=8,
        mlp_dim=2048,
    )
    mask_decoder = MaskDecoder(
        transformer_dim=embed_dim,
        transformer=transformer_decoder,
        num_multimask_outputs=3,
        iou_head_depth=3,
        iou_head_hidden_dim=256,
    )

    # Assemble full model
    model = PointCloudSAM(
        pc_encoder=pc_encoder,
        mask_encoder=mask_encoder,
        mask_decoder=mask_decoder,
        prompt_iters=prompt_iters,
    )

    if use_fused_layernorm:
        if replace_with_fused_layernorm is None:
            warnings.warn("apex FusedLayerNorm requested but not available. Skipping.")
        else:
            model = replace_with_fused_layernorm(model)

    model = model.to(device)
    return model


class PointSAM:
    """User-friendly wrapper around PointCloudSAM for single-file inference.

    Typical usage:
        >>> psam = PointSAM.from_pretrained("cuda")
        >>> coords, rgb, original = load_pointcloud("scene.ply")
        >>> mask = psam.predict(coords, rgb, prompt_point=[0.5, 0.1, -0.2])
    """

    def __init__(
        self,
        model: PointCloudSAM,
        device: str = "cuda",
        variant: str = "large",
    ):
        self.model = model
        self.device = device
        self.variant = variant
        self._pc_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None

    @classmethod
    def from_pretrained(
        cls,
        checkpoint_path: Optional[str] = None,
        variant: str = "large",
        device: str = "cuda",
        use_fused_layernorm: bool = False,
    ) -> "PointSAM":
        """Load a Point-SAM model from a local or Hub checkpoint.

        Args:
            checkpoint_path: Local path to a .safetensors checkpoint.
                If None, the model is initialized with random weights.
            variant: "large" or "giant".
            device: torch device.
            use_fused_layernorm: Whether to use apex FusedLayerNorm.

        Returns:
            PointSAM wrapper ready for inference.
        """
        model = build_point_sam(
            variant=variant,
            device=device,
            use_fused_layernorm=use_fused_layernorm,
        )
        if checkpoint_path is not None:
            load_safetensors_model(model, checkpoint_path)
            print(f"Loaded checkpoint from {checkpoint_path}")
        else:
            warnings.warn(
                "No checkpoint provided — model weights are randomly initialized!"
            )
        model.eval()
        return cls(model=model, device=device, variant=variant)

    def set_pointcloud(
        self,
        coords: Union[np.ndarray, torch.Tensor],
        rgb: Union[np.ndarray, torch.Tensor],
    ):
        """Cache a point cloud for repeated segmentation queries.

        This precomputes the encoder embeddings so that subsequent `predict`
        calls with different prompt points are much faster.

        Args:
            coords: [N, 3] coordinates (normalized to [-1, 1]).
            rgb: [N, 3] colors in [0, 255].
        """
        if isinstance(coords, np.ndarray):
            coords = torch.from_numpy(coords).float()
        if isinstance(rgb, np.ndarray):
            rgb = torch.from_numpy(rgb).float()

        # Ensure batch dim and normalize colors
        if coords.dim() == 2:
            coords = coords.unsqueeze(0)  # [1, N, 3]
        if rgb.dim() == 2:
            rgb = rgb.unsqueeze(0)  # [1, N, 3]

        if rgb.max() > 1.0 + 1e-6:
            rgb = rgb / 255.0

        coords = coords.to(self.device)
        rgb = rgb.to(self.device)

        self._pc_cache = (coords, rgb)

    def predict(
        self,
        coords: Union[np.ndarray, torch.Tensor],
        rgb: Union[np.ndarray, torch.Tensor],
        prompt_point: Union[list, tuple, np.ndarray, torch.Tensor],
        prompt_label: int = 1,
        multimask_output: bool = True,
        return_logits: bool = False,
    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
        """Run segmentation on a point cloud given a prompt point.

        Args:
            coords: [N, 3] normalized coordinates or a cached point cloud.
            rgb: [N, 3] colors in [0, 255]. Ignored if coords was cached via
                `set_pointcloud`.
            prompt_point: [x, y, z] coordinate of the prompt. Must be in the same
                normalized space as `coords` (i.e., [-1, 1] if you used the default
                `load_pointcloud` normalization).
            prompt_label: 1 for foreground (positive), 0 for background (negative).
            multimask_output: If True, return 3 mask candidates + IoU scores.
                If False, return a single mask.
            return_logits: If True, return raw logits instead of a boolean mask.

        Returns:
            If multimask_output=False:
                mask: [N] boolean array (or float logits if return_logits=True).
            If multimask_output=True:
                masks: [3, N] boolean array of candidate masks.
                iou_preds: [3] IoU confidence scores for each candidate.
        """
        # Use cached point cloud if available and coords wasn't passed fresh
        if self._pc_cache is not None and coords is None:
            coords, rgb = self._pc_cache
        else:
            if isinstance(coords, np.ndarray):
                coords = torch.from_numpy(coords).float()
            if isinstance(rgb, np.ndarray):
                rgb = torch.from_numpy(rgb).float()
            if coords.dim() == 2:
                coords = coords.unsqueeze(0)
            if rgb.dim() == 2:
                rgb = rgb.unsqueeze(0)
            if rgb.max() > 1.0 + 1e-6:
                rgb = rgb / 255.0
            coords = coords.to(self.device)
            rgb = rgb.to(self.device)

        # Prepare prompt
        if isinstance(prompt_point, (list, tuple)):
            prompt_point = np.array(prompt_point, dtype=np.float32)
        if isinstance(prompt_point, np.ndarray):
            prompt_point = torch.from_numpy(prompt_point).float()
        if prompt_point.dim() == 1:
            prompt_point = prompt_point.unsqueeze(0).unsqueeze(0)  # [1, 1, 3]
        prompt_point = prompt_point.to(self.device)

        prompt_labels = torch.tensor([[prompt_label]], dtype=torch.long, device=self.device)

        with torch.no_grad():
            masks, iou_preds = self.model.predict_masks(
                coords,
                rgb,
                prompt_point,
                prompt_labels,
                prompt_masks=None,
                multimask_output=multimask_output,
            )

        # masks: [1, num_outputs, N]
        # iou_preds: [1, num_outputs]
        masks = masks[0]  # [num_outputs, N]
        iou_preds = iou_preds[0]  # [num_outputs]

        if not multimask_output:
            mask = masks[0]
            if return_logits:
                return mask.cpu().numpy()
            return (mask > 0).cpu().numpy()

        if return_logits:
            return masks.cpu().numpy(), iou_preds.cpu().numpy()
        return (masks > 0).cpu().numpy(), iou_preds.cpu().numpy()

    @property
    def num_points(self) -> int:
        """Number of points in the cached point cloud, or 0 if none."""
        if self._pc_cache is None:
            return 0
        return self._pc_cache[0].shape[1]

    def adjust_patch_params(self, num_groups: int, group_size: int):
        """Dynamically adjust the number of patches and patch size.

        Call this when working with very large point clouds (e.g. > 100k points)
        to avoid OOM.  The authors suggest num_groups=2048, group_size=256 for
        clouds with > 100k points.
        """
        self.model.pc_encoder.patch_embed.grouper.num_groups = num_groups
        self.model.pc_encoder.patch_embed.grouper.group_size = group_size