sdf-sampler 0.1.0__py3-none-any.whl

sdf_sampler/__init__.py

@@ -0,0 +1,78 @@
+# ABOUTME: sdf-sampler public API exports
+# ABOUTME: Standalone SDF training data generation from point clouds
+
+"""
+sdf-sampler: Auto-analysis and sampling of point clouds for SDF training.
+
+Example usage:
+    from sdf_sampler import SDFAnalyzer, SDFSampler, load_point_cloud
+
+    # Load point cloud
+    xyz, normals = load_point_cloud("scan.ply")
+
+    # Auto-analyze to detect EMPTY/SOLID regions
+    analyzer = SDFAnalyzer()
+    result = analyzer.analyze(xyz=xyz, normals=normals)
+
+    # Generate training samples
+    sampler = SDFSampler()
+    samples = sampler.generate(
+        xyz=xyz,
+        normals=normals,
+        constraints=result.constraints,
+    )
+
+    # Export to parquet
+    sampler.export_parquet(samples, "training_data.parquet")
+"""
+
+from sdf_sampler.analyzer import SDFAnalyzer
+from sdf_sampler.config import AnalyzerConfig, SamplerConfig
+from sdf_sampler.io import export_parquet, load_point_cloud
+from sdf_sampler.models import (
+    AlgorithmType,
+    AnalysisResult,
+    BoxConstraint,
+    BrushStrokeConstraint,
+    Constraint,
+    HalfspaceConstraint,
+    PocketConstraint,
+    RayCarveConstraint,
+    SamplePointConstraint,
+    SamplingStrategy,
+    SeedPropagationConstraint,
+    SignConvention,
+    SphereConstraint,
+    TrainingSample,
+)
+from sdf_sampler.sampler import SDFSampler
+
+__version__ = "0.1.0"
+
+__all__ = [
+    # Main classes
+    "SDFAnalyzer",
+    "SDFSampler",
+    # Config
+    "AnalyzerConfig",
+    "SamplerConfig",
+    # I/O
+    "load_point_cloud",
+    "export_parquet",
+    # Models
+    "SignConvention",
+    "AlgorithmType",
+    "SamplingStrategy",
+    "AnalysisResult",
+    "TrainingSample",
+    # Constraints
+    "Constraint",
+    "BoxConstraint",
+    "SphereConstraint",
+    "HalfspaceConstraint",
+    "BrushStrokeConstraint",
+    "SeedPropagationConstraint",
+    "RayCarveConstraint",
+    "PocketConstraint",
+    "SamplePointConstraint",
+]

sdf_sampler/algorithms/__init__.py

@@ -0,0 +1,19 @@
+# ABOUTME: Algorithm module exports
+# ABOUTME: Provides analysis algorithms for SDF region detection
+
+from sdf_sampler.algorithms.flood_fill import flood_fill_empty_regions
+from sdf_sampler.algorithms.normal_idw import generate_idw_normal_samples
+from sdf_sampler.algorithms.normal_offset import generate_normal_offset_boxes
+from sdf_sampler.algorithms.pocket import detect_pockets
+from sdf_sampler.algorithms.voxel_grid import build_voxel_grid, compute_hull_mask
+from sdf_sampler.algorithms.voxel_regions import generate_voxel_region_constraints
+
+__all__ = [
+    "flood_fill_empty_regions",
+    "generate_voxel_region_constraints",
+    "generate_normal_offset_boxes",
+    "generate_idw_normal_samples",
+    "detect_pockets",
+    "build_voxel_grid",
+    "compute_hull_mask",
+]

sdf_sampler/algorithms/flood_fill.py

@@ -0,0 +1,233 @@
+# ABOUTME: Flood fill algorithm for EMPTY region detection
+# ABOUTME: Uses ray propagation from sky to identify exterior empty space
+
+import numpy as np
+from scipy.spatial import KDTree
+
+from sdf_sampler.algorithms.voxel_grid import (
+    build_voxel_grid,
+    compute_hull_mask,
+    greedy_2d_mesh,
+    ray_propagation_with_bounces,
+)
+from sdf_sampler.config import AutoAnalysisOptions
+from sdf_sampler.models.analysis import AlgorithmType, GeneratedConstraint
+from sdf_sampler.models.constraints import SignConvention
+
+
+def flood_fill_empty_regions(
+    xyz: np.ndarray,
+    normals: np.ndarray | None,
+    options: AutoAnalysisOptions,
+) -> list[GeneratedConstraint]:
+    """Generate EMPTY constraints using ray propagation with bouncing.
+
+    Uses the ray model:
+    1. EMPTY rays shine down from +Z (sky)
+    2. Rays bounce to fill occluded areas (trenches, overhangs)
+    3. Output depends on flood_fill_output option
+
+    Args:
+        xyz: Point cloud positions (N, 3)
+        normals: Point normals (N, 3) or None
+        options: Algorithm options
+
+    Returns:
+        List of GeneratedConstraint objects
+    """
+    constraints: list[GeneratedConstraint] = []
+
+    grid_result = build_voxel_grid(xyz, options.min_gap_size, options.max_grid_dim)
+    if grid_result is None:
+        return constraints
+
+    occupied, bbox_min, voxel_size, grid_shape = grid_result
+    _nx, _ny, nz = grid_shape
+
+    inside_hull = compute_hull_mask(xyz, bbox_min, voxel_size, grid_shape)
+
+    empty_mask, _ = ray_propagation_with_bounces(
+        occupied, grid_shape, inside_hull, options.cone_angle
+    )
+
+    output_mode = options.flood_fill_output.lower()
+
+    if output_mode in ("samples", "both"):
+        sample_constraints = _generate_samples_from_mask(
+            empty_mask,
+            bbox_min,
+            voxel_size,
+            xyz,
+            options.flood_fill_sample_count,
+            SignConvention.EMPTY,
+            AlgorithmType.FLOOD_FILL,
+        )
+        constraints.extend(sample_constraints)
+
+    if output_mode in ("boxes", "both"):
+        box_constraints = _generate_boxes_from_mask(
+            empty_mask,
+            bbox_min,
+            voxel_size,
+            nz,
+            options,
+            SignConvention.EMPTY,
+            AlgorithmType.FLOOD_FILL,
+        )
+        constraints.extend(box_constraints)
+
+    return constraints
+
+
+def _generate_samples_from_mask(
+    mask: np.ndarray,
+    bbox_min: np.ndarray,
+    voxel_size: float,
+    xyz: np.ndarray,
+    n_samples: int,
+    sign: SignConvention,
+    algorithm: AlgorithmType,
+) -> list[GeneratedConstraint]:
+    """Generate sample_point constraints from a voxel mask.
+
+    Uses inverse-square distance weighting: more samples near the surface,
+    fewer samples far away.
+    """
+    constraints: list[GeneratedConstraint] = []
+
+    marked_indices = np.argwhere(mask)
+    if len(marked_indices) == 0:
+        return constraints
+
+    tree = KDTree(xyz)
+    rng = np.random.default_rng(42)
+
+    voxel_centers = bbox_min + (marked_indices + 0.5) * voxel_size
+    distances, _ = tree.query(voxel_centers, k=1)
+
+    epsilon = voxel_size * 0.1
+    weights = 1.0 / (distances + epsilon) ** 2
+    weights = weights / weights.sum()
+
+    sample_indices = rng.choice(len(marked_indices), size=n_samples, replace=True, p=weights)
+
+    for idx in sample_indices:
+        voxel_ijk = marked_indices[idx]
+        offset = rng.uniform(0, 1, 3)
+        world_pos = bbox_min + (voxel_ijk + offset) * voxel_size
+
+        dist, _ = tree.query(world_pos, k=1)
+        signed_dist = float(dist) if sign == SignConvention.EMPTY else -float(dist)
+
+        constraints.append(
+            GeneratedConstraint(
+                constraint={
+                    "type": "sample_point",
+                    "sign": sign.value,
+                    "position": tuple(world_pos.tolist()),
+                    "distance": signed_dist,
+                },
+                algorithm=algorithm,
+                confidence=0.8,
+                description=f"Voxel sample at d={signed_dist:.3f}m",
+            )
+        )
+
+    return constraints
+
+
+def _generate_boxes_from_mask(
+    empty_mask: np.ndarray,
+    bbox_min: np.ndarray,
+    voxel_size: float,
+    nz: int,
+    options: AutoAnalysisOptions,
+    sign: SignConvention,
+    algorithm: AlgorithmType,
+) -> list[GeneratedConstraint]:
+    """Generate axis-aligned box constraints from a voxel mask using greedy meshing."""
+    constraints: list[GeneratedConstraint] = []
+
+    all_boxes: list[tuple[int, int, int, int, int, int]] = []
+
+    for iz in range(nz):
+        slice_2d = empty_mask[:, :, iz]
+        if not slice_2d.any():
+            continue
+
+        rectangles = greedy_2d_mesh(slice_2d)
+        for x_min, x_max, y_min, y_max in rectangles:
+            all_boxes.append((iz, iz + 1, x_min, x_max, y_min, y_max))
+
+    if not all_boxes:
+        return constraints
+
+    all_boxes.sort(key=lambda b: (b[2], b[3], b[4], b[5], b[0]))
+
+    merged_boxes: list[tuple[int, int, int, int, int, int]] = []
+    current = all_boxes[0]
+
+    for box in all_boxes[1:]:
+        z_start, z_end, x_min, x_max, y_min, y_max = box
+        cz_start, cz_end, cx_min, cx_max, cy_min, cy_max = current
+
+        if (
+            x_min == cx_min
+            and x_max == cx_max
+            and y_min == cy_min
+            and y_max == cy_max
+            and z_start == cz_end
+        ):
+            current = (cz_start, z_end, cx_min, cx_max, cy_min, cy_max)
+        else:
+            merged_boxes.append(current)
+            current = box
+
+    merged_boxes.append(current)
+
+    min_extent = 3
+    merged_boxes = [
+        b
+        for b in merged_boxes
+        if (b[1] - b[0]) >= min_extent
+        and (b[3] - b[2]) >= min_extent
+        and (b[5] - b[4]) >= min_extent
+    ]
+
+    if not merged_boxes:
+        return constraints
+
+    max_boxes = options.max_boxes
+    if len(merged_boxes) > max_boxes:
+        merged_boxes.sort(
+            key=lambda b: (b[1] - b[0]) * (b[3] - b[2]) * (b[5] - b[4]),
+            reverse=True,
+        )
+        merged_boxes = merged_boxes[:max_boxes]
+
+    for z_start, z_end, x_min, x_max, y_min, y_max in merged_boxes:
+        world_min = bbox_min + np.array([x_min, y_min, z_start]) * voxel_size
+        world_max = bbox_min + np.array([x_max, y_max, z_end]) * voxel_size
+
+        center = (world_min + world_max) / 2
+        half_extents = (world_max - world_min) / 2
+
+        box_constraint = {
+            "type": "box",
+            "sign": sign.value,
+            "center": tuple(center.tolist()),
+            "half_extents": tuple(half_extents.tolist()),
+        }
+
+        volume = float(np.prod(half_extents * 2))
+        n_voxels = (z_end - z_start) * (x_max - x_min) * (y_max - y_min)
+        constraints.append(
+            GeneratedConstraint(
+                constraint=box_constraint,
+                algorithm=algorithm,
+                confidence=0.85,
+                description=f"Sky-reachable region ({n_voxels} voxels, {volume:.2f}m³)",
+            )
+        )
+
+    return constraints

sdf_sampler/algorithms/normal_idw.py

@@ -0,0 +1,99 @@
+# ABOUTME: IDW normal sampling algorithm
+# ABOUTME: Generates sample constraints along normals with inverse distance weighting
+
+import numpy as np
+
+from sdf_sampler.algorithms.normal_offset import _farthest_point_sample
+from sdf_sampler.config import AutoAnalysisOptions
+from sdf_sampler.models.analysis import AlgorithmType, GeneratedConstraint
+
+
+def generate_idw_normal_samples(
+    xyz: np.ndarray,
+    normals: np.ndarray | None,
+    options: AutoAnalysisOptions,
+) -> list[GeneratedConstraint]:
+    """Generate sample constraints along normals with inverse distance weighting.
+
+    Creates point samples at varying distances along surface normals, with
+    more samples concentrated near the surface (IDW = 1/distance^power).
+
+    Args:
+        xyz: Point cloud positions (N, 3)
+        normals: Point normals (N, 3) - required for this algorithm
+        options: Algorithm options
+
+    Returns:
+        List of GeneratedConstraint objects
+    """
+    constraints: list[GeneratedConstraint] = []
+
+    if normals is None or len(normals) != len(xyz):
+        return constraints
+
+    oriented_normals = _orient_normals_outward(xyz, normals)
+
+    n_surface_pts = min(options.idw_sample_count // 10, len(xyz))
+    if n_surface_pts < 1:
+        return constraints
+
+    surface_indices = _farthest_point_sample(xyz, n_surface_pts)
+    samples_per_point = options.idw_sample_count // len(surface_indices)
+
+    if samples_per_point < 1:
+        samples_per_point = 1
+
+    rng = np.random.default_rng(42)
+
+    for idx in surface_indices:
+        point = xyz[idx]
+        normal = oriented_normals[idx]
+        normal_norm = np.linalg.norm(normal)
+        if normal_norm < 0.1:
+            continue
+        normal = normal / normal_norm
+
+        # Generate distances with IDW distribution
+        u = rng.random(samples_per_point)
+        distances = options.idw_max_distance * (1 - u ** (1 / options.idw_power))
+
+        for dist in distances:
+            sign = rng.choice([-1, 1])
+            sample_pos = point + sign * dist * normal
+            sample_sign = "empty" if sign > 0 else "solid"
+
+            constraints.append(
+                GeneratedConstraint(
+                    constraint={
+                        "type": "sample_point",
+                        "sign": sample_sign,
+                        "position": tuple(sample_pos.tolist()),
+                        "distance": float(sign * dist),
+                    },
+                    algorithm=AlgorithmType.NORMAL_IDW,
+                    confidence=0.8,
+                    description=f"IDW sample at d={sign * dist:.3f}m",
+                )
+            )
+
+    return constraints
+
+
+def _orient_normals_outward(xyz: np.ndarray, normals: np.ndarray) -> np.ndarray:
+    """Orient normals to point outward using a viewpoint heuristic.
+
+    Uses assumption that viewpoint is above the scene (outdoor scenes).
+    """
+    centroid = xyz.mean(axis=0)
+    z_range = xyz[:, 2].max() - xyz[:, 2].min()
+    viewpoint = centroid.copy()
+    viewpoint[2] = xyz[:, 2].max() + z_range * 0.5
+
+    to_viewpoint = viewpoint - xyz
+    dot_products = np.sum(normals * to_viewpoint, axis=1)
+
+    oriented = normals.copy()
+    flip_mask = dot_products < 0
+    oriented[flip_mask] = -oriented[flip_mask]
+
+    return oriented

sdf_sampler/algorithms/normal_offset.py

@@ -0,0 +1,111 @@
+# ABOUTME: Normal offset algorithm for surface-relative constraints
+# ABOUTME: Generates paired SOLID/EMPTY boxes offset along surface normals
+
+import numpy as np
+from scipy.spatial import KDTree
+
+from sdf_sampler.algorithms.voxel_grid import estimate_mean_spacing
+from sdf_sampler.config import AutoAnalysisOptions
+from sdf_sampler.models.analysis import AlgorithmType, GeneratedConstraint
+from sdf_sampler.models.constraints import SignConvention
+
+
+def generate_normal_offset_boxes(
+    xyz: np.ndarray,
+    normals: np.ndarray | None,
+    options: AutoAnalysisOptions,
+) -> list[GeneratedConstraint]:
+    """Generate paired SOLID/EMPTY boxes offset along surface normals.
+
+    For a surface point with normal N:
+    - Box offset in +N direction (outward) -> EMPTY
+    - Box offset in -N direction (inward) -> SOLID
+
+    Args:
+        xyz: Point cloud positions (N, 3)
+        normals: Point normals (N, 3) - required for this algorithm
+        options: Algorithm options
+
+    Returns:
+        List of GeneratedConstraint objects
+    """
+    constraints: list[GeneratedConstraint] = []
+
+    if normals is None or len(normals) != len(xyz):
+        return constraints
+
+    tree = KDTree(xyz)
+    mean_spacing = estimate_mean_spacing(xyz, tree)
+
+    sample_indices = _farthest_point_sample(xyz, options.normal_offset_pairs)
+    offset_distance = mean_spacing * 3
+    box_size = mean_spacing * 2.5
+
+    for idx in sample_indices:
+        point = xyz[idx]
+        normal = normals[idx]
+        normal_norm = np.linalg.norm(normal)
+
+        if normal_norm < 0.1:
+            continue
+
+        normal = normal / normal_norm
+
+        # EMPTY box in +normal direction (outward)
+        empty_center = point + normal * offset_distance
+        box_constraint_empty = {
+            "type": "box",
+            "sign": SignConvention.EMPTY.value,
+            "center": tuple(empty_center.tolist()),
+            "half_extents": (box_size, box_size, box_size),
+        }
+
+        constraints.append(
+            GeneratedConstraint(
+                constraint=box_constraint_empty,
+                algorithm=AlgorithmType.NORMAL_OFFSET,
+                confidence=0.75,
+                description=f"Exterior offset from surface at ({point[0]:.2f}, {point[1]:.2f}, {point[2]:.2f})",
+            )
+        )
+
+        # SOLID box in -normal direction (inward)
+        solid_center = point - normal * offset_distance
+        box_constraint_solid = {
+            "type": "box",
+            "sign": SignConvention.SOLID.value,
+            "center": tuple(solid_center.tolist()),
+            "half_extents": (box_size, box_size, box_size),
+        }
+
+        constraints.append(
+            GeneratedConstraint(
+                constraint=box_constraint_solid,
+                algorithm=AlgorithmType.NORMAL_OFFSET,
+                confidence=0.75,
+                description=f"Interior offset from surface at ({point[0]:.2f}, {point[1]:.2f}, {point[2]:.2f})",
+            )
+        )
+
+    return constraints
+
+
+def _farthest_point_sample(xyz: np.ndarray, n_samples: int) -> list[int]:
+    """Select well-distributed points using farthest point sampling."""
+    n_points = len(xyz)
+    if n_samples >= n_points:
+        return list(range(n_points))
+
+    rng = np.random.default_rng(42)
+    selected: list[int] = [int(rng.integers(n_points))]
+    min_distances = np.full(n_points, np.inf)
+
+    for _ in range(n_samples - 1):
+        last_selected = xyz[selected[-1]]
+        distances = np.linalg.norm(xyz - last_selected, axis=1)
+        min_distances = np.minimum(min_distances, distances)
+        min_distances[selected] = -1
+        next_idx = int(np.argmax(min_distances))
+        selected.append(next_idx)
+
+    return selected

sdf_sampler/algorithms/pocket.py

@@ -0,0 +1,146 @@
+# ABOUTME: Pocket detection algorithm for cavity identification
+# ABOUTME: Uses voxel flood fill to find disconnected interior cavities
+
+from dataclasses import dataclass
+from enum import IntEnum
+
+import numpy as np
+from scipy import ndimage
+
+from sdf_sampler.config import AnalyzerConfig
+from sdf_sampler.models.analysis import AlgorithmType, GeneratedConstraint
+from sdf_sampler.models.constraints import SignConvention
+
+
+class VoxelState(IntEnum):
+    """State of a voxel in the occupancy grid."""
+
+    EMPTY = 0
+    OCCUPIED = 1
+    OUTSIDE = 2
+
+
+@dataclass
+class PocketInfo:
+    """Information about a detected pocket."""
+
+    pocket_id: int
+    voxel_count: int
+    centroid: tuple[float, float, float]
+    bounds_low: tuple[float, float, float]
+    bounds_high: tuple[float, float, float]
+    volume_estimate: float
+    voxel_coords: np.ndarray  # World coordinates of voxel centers
+
+
+def detect_pockets(
+    xyz: np.ndarray,
+    config: AnalyzerConfig,
+) -> list[GeneratedConstraint]:
+    """Detect pocket cavities in point cloud using voxel analysis.
+
+    Pockets are disconnected interior voids that should be marked as SOLID.
+
+    Args:
+        xyz: Point cloud positions (N, 3)
+        config: Analyzer configuration
+
+    Returns:
+        List of PocketConstraint GeneratedConstraints
+    """
+    constraints: list[GeneratedConstraint] = []
+
+    if len(xyz) < 10:
+        return constraints
+
+    # Compute bounds
+    bounds_low = xyz.min(axis=0)
+    bounds_high = xyz.max(axis=0)
+
+    # Compute voxel size
+    extent = bounds_high - bounds_low
+    longest_axis = np.max(extent)
+    voxel_size = longest_axis / config.pocket_voxel_target
+    voxel_size = max(voxel_size, config.pocket_min_voxel_size)
+
+    # Compute grid resolution - ensure at least 1 in each dimension
+    resolution = np.ceil(extent / voxel_size).astype(int)
+    resolution = np.maximum(resolution, 1)  # Avoid zero-size dimensions
+    resolution = np.minimum(resolution, config.pocket_max_voxels_per_axis)
+
+    # Build occupancy grid
+    grid = np.full(tuple(resolution), VoxelState.EMPTY, dtype=np.uint8)
+
+    voxel_indices = ((xyz - bounds_low) / voxel_size).astype(int)
+    voxel_indices = np.clip(voxel_indices, 0, resolution - 1)
+
+    # Mark occupied with dilation
+    if config.pocket_occupancy_dilation == 0:
+        grid[voxel_indices[:, 0], voxel_indices[:, 1], voxel_indices[:, 2]] = VoxelState.OCCUPIED
+    else:
+        occupied_mask = np.zeros_like(grid, dtype=bool)
+        occupied_mask[voxel_indices[:, 0], voxel_indices[:, 1], voxel_indices[:, 2]] = True
+        struct = ndimage.generate_binary_structure(3, 1)
+        dilated = ndimage.binary_dilation(
+            occupied_mask, structure=struct, iterations=config.pocket_occupancy_dilation
+        )
+        grid[dilated] = VoxelState.OCCUPIED
+
+    # Flood fill from boundary to mark outside
+    boundary_mask = np.zeros_like(grid, dtype=bool)
+    boundary_mask[0, :, :] = True
+    boundary_mask[-1, :, :] = True
+    boundary_mask[:, 0, :] = True
+    boundary_mask[:, -1, :] = True
+    boundary_mask[:, :, 0] = True
+    boundary_mask[:, :, -1] = True
+
+    seed = boundary_mask & (grid == VoxelState.EMPTY)
+    traversable = grid == VoxelState.EMPTY
+
+    struct = ndimage.generate_binary_structure(3, 1)
+    outside = ndimage.binary_dilation(seed, mask=traversable, iterations=-1, structure=struct)
+    grid[outside] = VoxelState.OUTSIDE
+
+    # Label remaining empty voxels as pockets
+    pocket_mask = grid == VoxelState.EMPTY
+    labeled, num_pockets = ndimage.label(pocket_mask, structure=struct)
+
+    # Extract pocket info
+    for pocket_id in range(1, num_pockets + 1):
+        voxel_count = int(np.sum(labeled == pocket_id))
+        if voxel_count < config.pocket_min_volume_voxels:
+            continue
+
+        mask = labeled == pocket_id
+        voxel_coords = np.argwhere(mask)
+        world_coords = voxel_coords * voxel_size + bounds_low + voxel_size / 2
+
+        centroid = tuple(world_coords.mean(axis=0).tolist())
+        pocket_bounds_low = tuple((voxel_coords.min(axis=0) * voxel_size + bounds_low).tolist())
+        pocket_bounds_high = tuple(
+            ((voxel_coords.max(axis=0) + 1) * voxel_size + bounds_low).tolist()
+        )
+        volume = voxel_count * (voxel_size**3)
+
+        pocket_constraint = {
+            "type": "pocket",
+            "sign": SignConvention.SOLID.value,
+            "pocket_id": pocket_id,
+            "voxel_count": voxel_count,
+            "centroid": centroid,
+            "bounds_low": pocket_bounds_low,
+            "bounds_high": pocket_bounds_high,
+            "volume_estimate": volume,
+        }
+
+        constraints.append(
+            GeneratedConstraint(
+                constraint=pocket_constraint,
+                algorithm=AlgorithmType.POCKET,
+                confidence=0.95,
+                description=f"Interior cavity at ({centroid[0]:.2f}, {centroid[1]:.2f}, {centroid[2]:.2f}), {voxel_count} voxels",
+            )
+        )
+
+    return constraints