fix: pick tools keep identity and place particles at their centre (v2.0)

pyproject.toml

@@ -19,7 +19,7 @@ classifiers = [
 ]
 dynamic = ["version"]
 dependencies = [
-    "copick>=2.0.0a1",
+    "copick>=2.0.0a3,<3",
     "click",
     "click-option-group",
     "numpy",

src/copick_utils/cli/clippicks.py

@@ -65,6 +65,11 @@ def clippicks(
     Exactly one reference must be supplied via `--ref-mesh`, `--ref-seg`, or
     `--ref-tomogram`. Distances are measured in angstroms.
 
+    Each pick is measured at its particle centre (its location plus the shift stored in its
+    transform). Kept picks are written unchanged and in their original order, so their
+    orientations, instance IDs (e.g. filament IDs) and scores survive. If no pick is kept, an
+    empty pick set is written.
+
     URI Format:
 
         \b

src/copick_utils/cli/picksin.py

@@ -50,6 +50,11 @@ def picksin(
     whose coordinates fall inside it are written to the output. This is the inverse of
     `copick logical picksout`, which keeps the picks that fall outside the volume instead.
 
+    Each pick is tested at its particle centre (its location plus the shift stored in its
+    transform). Kept picks are written unchanged and in their original order, so their
+    orientations, instance IDs (e.g. filament IDs) and scores survive. If no pick is kept, an
+    empty pick set is written.
+
     URI Format:
 
         \b

src/copick_utils/cli/picksout.py

@@ -52,6 +52,11 @@ def picksout(
     whose coordinates fall outside it are written to the output. This is the inverse of
     `copick logical picksin`, which keeps the picks that fall inside the volume instead.
 
+    Each pick is tested at its particle centre (its location plus the shift stored in its
+    transform). Kept picks are written unchanged and in their original order, so their
+    orientations, instance IDs (e.g. filament IDs) and scores survive. If no pick is kept, an
+    empty pick set is written.
+
     URI Format:
 
         \b

src/copick_utils/converters/ellipsoid_from_picks.py

@@ -11,6 +11,7 @@
     validate_points,
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.util.picks import pick_centres
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickRun
@@ -317,9 +318,10 @@ def create_ellipsoid_from_points(cluster_points):
 
 
 def ellipsoid_from_picks_standard(picks, run, object_name, session_id, user_id, voxel_spacing=None, **kwargs):
-    """Adapt ellipsoid_from_picks to the lazy-batch contract: take a CopickPicks and extract its points."""
-    pos, _ = picks.numpy()
-    if pos is None or len(pos) == 0:
+    """Adapt ellipsoid_from_picks to the lazy-batch contract: take a CopickPicks and extract its particle centres
+    (``location`` plus the transform's shift)."""
+    pos = pick_centres(picks)
+    if len(pos) == 0:
         logger.warning("Could not load pick data")
         return None
     return ellipsoid_from_picks(pos, run, object_name, session_id, user_id, **kwargs)

src/copick_utils/converters/mesh_from_picks.py

@@ -10,6 +10,7 @@
     validate_points,
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.util.picks import pick_centres
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickRun
@@ -189,9 +190,10 @@ def create_mesh_from_points(cluster_points):
 
 # Lazy batch converter for new architecture
 def mesh_from_picks_standard(picks, run, object_name, session_id, user_id, voxel_spacing=None, **kwargs):
-    """Adapt mesh_from_picks to the lazy-batch contract: take a CopickPicks and extract its points."""
-    pos, _ = picks.numpy()
-    if pos is None or len(pos) == 0:
+    """Adapt mesh_from_picks to the lazy-batch contract: take a CopickPicks and extract its particle centres
+    (``location`` plus the transform's shift)."""
+    pos = pick_centres(picks)
+    if len(pos) == 0:
         logger.warning("Could not load pick data")
         return None
     return mesh_from_picks(pos, run, object_name, session_id, user_id, **kwargs)

src/copick_utils/converters/plane_from_picks.py

@@ -11,6 +11,7 @@
     validate_points,
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.util.picks import pick_centres
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickRun
@@ -232,9 +233,10 @@ def create_plane_from_points(cluster_points):
 
 # Lazy batch converter for new architecture
 def plane_from_picks_standard(picks, run, object_name, session_id, user_id, voxel_spacing=None, **kwargs):
-    """Adapt plane_from_picks to the lazy-batch contract: take a CopickPicks and extract its points."""
-    pos, _ = picks.numpy()
-    if pos is None or len(pos) == 0:
+    """Adapt plane_from_picks to the lazy-batch contract: take a CopickPicks and extract its particle centres
+    (``location`` plus the transform's shift)."""
+    pos = pick_centres(picks)
+    if len(pos) == 0:
         logger.warning("Could not load pick data")
         return None
     return plane_from_picks(pos, run, object_name, session_id, user_id, **kwargs)

src/copick_utils/converters/segmentation_from_picks.py

@@ -5,6 +5,7 @@
 
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
 from copick_utils.io.zarr import get_level_array
+from copick_utils.util.picks import pick_centres
 
 if TYPE_CHECKING:
     from copick.models import CopickObject, CopickPicks, CopickRun, CopickSegmentation
@@ -54,14 +55,10 @@ def get_relative_target_coordinates(center, delta, shape):
     radius_voxel = max(radius / voxel_spacing, 1)
     delta = int(np.ceil(radius_voxel))
 
-    # Paint each pick as a sphere
-    for point in pick.points:
-        # Convert the pick's location from angstroms to voxel units
-        cx, cy, cz = (
-            point.location.x / voxel_spacing,
-            point.location.y / voxel_spacing,
-            point.location.z / voxel_spacing,
-        )
+    # Paint each pick as a sphere at its particle centre (location plus the transform's shift)
+    for centre in pick_centres(pick):
+        # Convert the particle centre from angstroms to voxel units
+        cx, cy, cz = centre / voxel_spacing
 
         # Calculate subarray bounds
         xLow, xHigh = get_relative_target_coordinates(cx, delta, seg_volume.shape[2])

src/copick_utils/converters/sphere_from_picks.py

@@ -11,6 +11,7 @@
     validate_points,
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.util.picks import pick_centres
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickRun
@@ -287,9 +288,10 @@ def sphere_from_picks(
 
 # Lazy batch converter for new architecture
 def sphere_from_picks_standard(picks, run, object_name, session_id, user_id, voxel_spacing=None, **kwargs):
-    """Adapt sphere_from_picks to the lazy-batch contract: take a CopickPicks and extract its points."""
-    pos, _ = picks.numpy()
-    if pos is None or len(pos) == 0:
+    """Adapt sphere_from_picks to the lazy-batch contract: take a CopickPicks and extract its particle centres
+    (``location`` plus the transform's shift)."""
+    pos = pick_centres(picks)
+    if len(pos) == 0:
         logger.warning("Could not load pick data")
         return None
     return sphere_from_picks(pos, run, object_name, session_id, user_id, **kwargs)

src/copick_utils/converters/surface_from_picks.py

@@ -13,6 +13,7 @@
     validate_points,
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+from copick_utils.util.picks import pick_centres
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickRun
@@ -318,9 +319,10 @@ def create_surface_from_points(cluster_points):
 
 # Lazy batch converter for new architecture
 def surface_from_picks_standard(picks, run, object_name, session_id, user_id, voxel_spacing=None, **kwargs):
-    """Adapt surface_from_picks to the lazy-batch contract: take a CopickPicks and extract its points."""
-    pos, _ = picks.numpy()
-    if pos is None or len(pos) == 0:
+    """Adapt surface_from_picks to the lazy-batch contract: take a CopickPicks and extract its particle centres
+    (``location`` plus the transform's shift)."""
+    pos = pick_centres(picks)
+    if len(pos) == 0:
         logger.warning("Could not load pick data")
         return None
     return surface_from_picks(pos, run, object_name, session_id, user_id, **kwargs)

src/copick_utils/io/readers.py

@@ -1,6 +1,7 @@
-import numpy as np
 from copick.util.uri import resolve_copick_objects
 
+from copick_utils.util.picks import pick_centres
+
 
 def tomogram(run, voxel_size: float = 10, algorithm: str = "wbp", raise_error: bool = False, verbose=True):
     """Read a tomogram from a copick run as a NumPy array.
@@ -143,7 +144,8 @@ def coordinates(
     """Read pick coordinates from a copick run as an array of voxel indices.
 
     Looks up the picks for `name` (optionally filtered by `user_id` and
-    `session_id`) and returns their locations divided by `voxel_size`. If several
+    `session_id`) and returns their particle centres (location plus the transform's
+    shift) divided by `voxel_size`. If several
     pick sets match, the first is used; if none match, reports the available picks.
 
     Args:
@@ -194,19 +196,5 @@ def coordinates(
                 f"Defaulting to loading:\n {picks[0]}\n",
             )
 
-    points = picks[0].points
-
-    # Initialize an array to store the coordinates
-    nPoints = len(picks[0].points)  # Number of points retrieved
-    coordinates = np.zeros([len(picks[0].points), 3])  # Create an empty array to hold the (z, y, x) coordinates
-
-    # Iterate over all points and convert their locations to coordinates in voxel space
-    for ii in range(nPoints):
-        coordinates[ii,] = [
-            points[ii].location.z / voxel_size,  # Scale z-coordinate by voxel size
-            points[ii].location.y / voxel_size,  # Scale y-coordinate by voxel size
-            points[ii].location.x / voxel_size,
-        ]  # Scale x-coordinate by voxel size
-
-    # Return the array of coordinates
-    return coordinates
+    # Particle centres (location plus the transform's shift), converted to (z, y, x) voxel coordinates
+    return pick_centres(picks[0])[:, ::-1] / voxel_size

src/copick_utils/logical/distance_operations.py

@@ -11,6 +11,7 @@
 )
 from copick_utils.converters.lazy_converter import create_lazy_batch_converter
 from copick_utils.io.zarr import get_level_array
+from copick_utils.util.picks import point_centres, segmentation_label_volume, store_pick_subset, voxel_indices
 
 if TYPE_CHECKING:
     from copick.models import CopickMesh, CopickPicks, CopickRun, CopickSegmentation
@@ -585,6 +586,10 @@ def limit_picks_by_distance(
     """
     Limit picks to those within a certain distance of a reference surface.
 
+    A pick is measured at its particle centre (``location`` plus the transform's shift). Kept picks are written
+    unchanged and in their original order: transform, ``instance_id`` and ``score`` survive. If no pick passes, an
+    empty pick set is written.
+
     Args:
         picks: CopickPicks to limit
         run: CopickRun object
@@ -609,13 +614,14 @@ def limit_picks_by_distance(
                 "One of reference_mesh, reference_segmentation, or reference_tomogram_info must be provided",
             )
 
-        # Load pick data
-        points, transforms = picks.numpy()
-        if points is None or len(points) == 0:
+        # Load pick data; picks are tested at their particle centres (location plus the transform's shift)
+        points = picks.points or []
+        if len(points) == 0:
             logger.error("Could not load pick data")
             return None
 
-        pick_positions = points[:, :3]  # Use only x, y, z coordinates
+        pick_positions = point_centres(points)
+        within_or_beyond = "beyond" if invert else "within"
 
         # Handle tomogram boundary reference (direct distance calculation, no distance field)
         if reference_tomogram_info is not None:
@@ -626,27 +632,19 @@ def limit_picks_by_distance(
             # Apply invert logic: default keeps within distance, invert keeps beyond
             final_valid = pick_distances > max_distance if invert else pick_distances <= max_distance
 
-            if not np.any(final_valid):
-                within_or_beyond = "beyond" if invert else "within"
-                logger.warning(f"No picks {within_or_beyond} {max_distance} units of tomogram boundary")
-                return None
-
         # Handle segmentation or mesh reference (uses distance field)
         else:
             # Create distance field from reference
             if reference_segmentation is not None:
-                ref_seg_array = reference_segmentation.numpy()
+                ref_seg_array = segmentation_label_volume(reference_segmentation)
                 if ref_seg_array is None or ref_seg_array.size == 0:
                     logger.error("Could not load reference segmentation data")
                     return None
 
-                # Use reference segmentation's coordinate space
+                # The field is on the reference segmentation's own (z, y, x) grid; voxel i sits at i * spacing.
                 field_voxel_spacing = reference_segmentation.voxel_size
                 distance_field = _create_distance_field_from_segmentation(ref_seg_array, field_voxel_spacing)
-
-                # Convert pick coordinates to voxel indices in reference segmentation space
-                pick_voxel_coords = pick_positions / field_voxel_spacing
-                pick_voxel_indices = np.floor(pick_voxel_coords).astype(int)
+                pick_voxel_indices = voxel_indices(pick_positions, field_voxel_spacing)
 
             else:  # reference_mesh is not None
                 ref_mesh = reference_mesh.mesh
@@ -673,7 +671,7 @@ def limit_picks_by_distance(
                 field_size = coord_bounds[1] - coord_bounds[0]
                 field_shape = np.ceil(field_size / field_voxel_spacing).astype(int)
 
-                # Create distance field from mesh in this coordinate space.
+                # Create distance field from mesh in this coordinate space (axis order x, y, z).
                 # Pass origin=coord_bounds[0] so the field aligns with the pick-index origin below.
                 distance_field = _create_distance_field_from_mesh(
                     ref_mesh,
@@ -692,10 +690,6 @@ def limit_picks_by_distance(
                 axis=1,
             )
 
-            if not np.any(valid_picks):
-                logger.warning("No picks within distance field bounds")
-                return None
-
             # Get distances for valid picks
             valid_indices = pick_voxel_indices[valid_picks]
             pick_distances = distance_field[valid_indices[:, 0], valid_indices[:, 1], valid_indices[:, 2]]
@@ -707,22 +701,15 @@ def limit_picks_by_distance(
             final_valid = np.zeros(len(points), dtype=bool)
             final_valid[valid_picks] = distance_valid
 
-            if not np.any(final_valid):
-                within_or_beyond = "beyond" if invert else "within"
-                logger.warning(f"No picks {within_or_beyond} {max_distance} units of reference surface")
-                return None
-
-        # Filter picks
-        valid_points = points[final_valid]
-        valid_transforms = transforms[final_valid] if transforms is not None else None
+        if not np.any(final_valid):
+            logger.warning(
+                f"No picks {within_or_beyond} {max_distance} units of the reference; writing an empty pick set",
+            )
 
-        # Create output picks
-        output_picks = run.new_picks(object_name, session_id, user_id, exist_ok=True)
-        output_picks.from_numpy(positions=valid_points, transforms=valid_transforms)
-        output_picks.store()
+        # Kept picks are written unchanged and in order (transform, instance_id and score survive)
+        output_picks = store_pick_subset(points, final_valid, run, object_name, session_id, user_id)
 
-        stats = {"points_created": len(valid_points)}
-        within_or_beyond = "beyond" if invert else "within"
+        stats = {"points_created": int(np.count_nonzero(final_valid))}
         logger.info(f"Limited picks to {stats['points_created']} points {within_or_beyond} {max_distance} units")
         return output_picks, stats