copick-utils 0.6.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

copick_utils/converters/surface_from_picks.py

@@ -0,0 +1,337 @@
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
+
+import numpy as np
+import trimesh as tm
+from copick.util.log import get_logger
+from scipy.interpolate import Rbf, griddata
+from scipy.spatial import Delaunay
+from sklearn.decomposition import PCA
+
+from copick_utils.converters.converter_common import (
+    cluster,
+    create_batch_converter,
+    create_batch_worker,
+    store_mesh_with_stats,
+    validate_points,
+)
+from copick_utils.converters.lazy_converter import create_lazy_batch_converter
+
+if TYPE_CHECKING:
+    from copick.models import CopickMesh, CopickRun
+
+logger = get_logger(__name__)
+
+
+def fit_2d_surface_to_points(
+    points: np.ndarray,
+    method: str = "delaunay",
+    grid_resolution: int = 50,
+) -> tm.Trimesh:
+    """Fit a 2D surface to 3D points using different interpolation methods.
+
+    Args:
+        points: Nx3 array of points.
+        method: Surface fitting method ('delaunay', 'rbf', 'grid').
+        grid_resolution: Resolution for grid-based methods.
+
+    Returns:
+        Trimesh surface object.
+    """
+    if len(points) < 3:
+        raise ValueError("Need at least 3 points to fit a surface")
+
+    if method == "delaunay":
+        return delaunay_surface(points)
+    elif method == "rbf":
+        return rbf_surface(points, grid_resolution)
+    elif method == "grid":
+        return grid_surface(points, grid_resolution)
+    else:
+        raise ValueError(f"Unknown surface method: {method}")
+
+
+def delaunay_surface(points: np.ndarray) -> tm.Trimesh:
+    """Create a surface using Delaunay triangulation.
+
+    Args:
+        points: Nx3 array of points.
+
+    Returns:
+        Trimesh surface object.
+    """
+    # Find the best 2D projection plane using PCA
+    center = np.mean(points, axis=0)
+    centered_points = points - center
+
+    pca = PCA(n_components=3)
+    pca.fit(centered_points)
+
+    # Use first two principal components for 2D projection
+    projected_2d = pca.transform(centered_points)[:, :2]
+
+    # Create Delaunay triangulation in 2D
+    tri = Delaunay(projected_2d)
+
+    # Use original 3D points as vertices with 2D triangulation
+    return tm.Trimesh(vertices=points, faces=tri.simplices)
+
+
+def rbf_surface(points: np.ndarray, grid_resolution: int) -> tm.Trimesh:
+    """Create a surface using RBF (Radial Basis Function) interpolation.
+
+    Args:
+        points: Nx3 array of points.
+        grid_resolution: Resolution of the output grid.
+
+    Returns:
+        Trimesh surface object.
+    """
+    # Find the dominant plane using PCA
+    center = np.mean(points, axis=0)
+    centered_points = points - center
+
+    pca = PCA(n_components=3)
+    pca.fit(centered_points)
+
+    # Project points onto the first two principal components
+    projected_2d = pca.transform(centered_points)[:, :2]
+    heights = pca.transform(centered_points)[:, 2]  # Third component as height
+
+    # Create grid for interpolation
+    x_min, x_max = projected_2d[:, 0].min(), projected_2d[:, 0].max()
+    y_min, y_max = projected_2d[:, 1].min(), projected_2d[:, 1].max()
+
+    xi = np.linspace(x_min, x_max, grid_resolution)
+    yi = np.linspace(y_min, y_max, grid_resolution)
+    xi_grid, yi_grid = np.meshgrid(xi, yi)
+
+    # RBF interpolation
+    rbf = Rbf(projected_2d[:, 0], projected_2d[:, 1], heights, function="thin_plate")
+    zi_grid = rbf(xi_grid, yi_grid)
+
+    # Convert grid back to 3D coordinates
+    grid_points_2d = np.column_stack([xi_grid.flatten(), yi_grid.flatten(), zi_grid.flatten()])
+    grid_points_3d = pca.inverse_transform(grid_points_2d) + center
+
+    # Create triangulation for the grid
+    grid_points_3d.reshape((grid_resolution, grid_resolution, 3))
+    faces = []
+
+    for i in range(grid_resolution - 1):
+        for j in range(grid_resolution - 1):
+            # Two triangles per grid cell
+            v1 = i * grid_resolution + j
+            v2 = i * grid_resolution + (j + 1)
+            v3 = (i + 1) * grid_resolution + j
+            v4 = (i + 1) * grid_resolution + (j + 1)
+
+            faces.extend([[v1, v2, v3], [v2, v4, v3]])
+
+    return tm.Trimesh(vertices=grid_points_3d, faces=faces)
+
+
+def grid_surface(points: np.ndarray, grid_resolution: int) -> tm.Trimesh:
+    """Create a surface using grid-based interpolation.
+
+    Args:
+        points: Nx3 array of points.
+        grid_resolution: Resolution of the output grid.
+
+    Returns:
+        Trimesh surface object.
+    """
+    # Find bounding box
+    min_coords = np.min(points, axis=0)
+    max_coords = np.max(points, axis=0)
+
+    # Find the dimension with smallest range (likely the "height" dimension)
+    ranges = max_coords - min_coords
+    height_dim = np.argmin(ranges)
+
+    # Use other two dimensions for grid
+    other_dims = [i for i in range(3) if i != height_dim]
+
+    # Create grid
+    x_coords = np.linspace(min_coords[other_dims[0]], max_coords[other_dims[0]], grid_resolution)
+    y_coords = np.linspace(min_coords[other_dims[1]], max_coords[other_dims[1]], grid_resolution)
+    xi, yi = np.meshgrid(x_coords, y_coords)
+
+    # Interpolate height values
+    zi = griddata(
+        points[:, other_dims],
+        points[:, height_dim],
+        (xi, yi),
+        method="linear",
+        fill_value=np.mean(points[:, height_dim]),
+    )
+
+    # Build 3D vertices
+    vertices = np.zeros((grid_resolution * grid_resolution, 3))
+    vertices[:, other_dims[0]] = xi.flatten()
+    vertices[:, other_dims[1]] = yi.flatten()
+    vertices[:, height_dim] = zi.flatten()
+
+    # Create triangulation
+    faces = []
+    for i in range(grid_resolution - 1):
+        for j in range(grid_resolution - 1):
+            # Two triangles per grid cell
+            v1 = i * grid_resolution + j
+            v2 = i * grid_resolution + (j + 1)
+            v3 = (i + 1) * grid_resolution + j
+            v4 = (i + 1) * grid_resolution + (j + 1)
+
+            faces.extend([[v1, v2, v3], [v2, v4, v3]])
+
+    return tm.Trimesh(vertices=vertices, faces=faces)
+
+
+def surface_from_picks(
+    points: np.ndarray,
+    run: "CopickRun",
+    object_name: str,
+    session_id: str,
+    user_id: str,
+    surface_method: str = "delaunay",
+    grid_resolution: int = 50,
+    use_clustering: bool = False,
+    clustering_method: str = "dbscan",
+    clustering_params: Optional[Dict[str, Any]] = None,
+    all_clusters: bool = True,
+    individual_meshes: bool = False,
+    session_id_template: Optional[str] = None,
+) -> Optional[Tuple["CopickMesh", Dict[str, int]]]:
+    """Create surface mesh(es) from pick points.
+
+    Args:
+        points: Nx3 array of pick positions.
+        run: Copick run object.
+        object_name: Name of the mesh object.
+        session_id: Session ID for the mesh.
+        user_id: User ID for the mesh.
+        surface_method: Surface fitting method ('delaunay', 'rbf', 'grid').
+        grid_resolution: Resolution for grid-based methods.
+        use_clustering: Whether to cluster points first.
+        clustering_method: Clustering method ('dbscan', 'kmeans').
+        clustering_params: Parameters for clustering.
+            e.g.
+                - {'eps': 5.0, 'min_samples': 3} for DBSCAN
+                - {'n_clusters': 3} for KMeans
+        all_clusters: If True, use all clusters; if False, use only the largest cluster.
+        individual_meshes: If True, create separate mesh objects for each surface.
+        session_id_template: Template for individual mesh session IDs.
+
+    Returns:
+        Tuple of (CopickMesh object, stats dict) or None if creation failed.
+        Stats dict contains 'vertices_created' and 'faces_created' totals.
+    """
+    if not validate_points(points, 3, "surface"):
+        return None
+
+    if clustering_params is None:
+        clustering_params = {}
+
+    # Define surface creation function
+    def create_surface_from_points(cluster_points):
+        return fit_2d_surface_to_points(cluster_points, surface_method, grid_resolution)
+
+    # Handle clustering workflow with special surface logic
+    if use_clustering:
+        point_clusters = cluster(
+            points,
+            clustering_method,
+            min_points_per_cluster=3,  # Surfaces need at least 3 points
+            **clustering_params,
+        )
+
+        if not point_clusters:
+            logger.warning("No valid clusters found")
+            return None
+
+        logger.info(f"Found {len(point_clusters)} clusters")
+
+        if all_clusters and len(point_clusters) > 1:
+            if individual_meshes:
+                # Create separate mesh objects for each surface
+                created_meshes = []
+                total_vertices = 0
+                total_faces = 0
+
+                for i, cluster_points in enumerate(point_clusters):
+                    try:
+                        surface_mesh = create_surface_from_points(cluster_points)
+
+                        # Generate session ID using template if provided
+                        if session_id_template:
+                            surface_session_id = session_id_template.format(
+                                base_session_id=session_id,
+                                instance_id=i,
+                            )
+                        else:
+                            surface_session_id = f"{session_id}-{i:03d}"
+
+                        copick_mesh = run.new_mesh(object_name, surface_session_id, user_id, exist_ok=True)
+                        copick_mesh.mesh = surface_mesh
+                        copick_mesh.store()
+                        created_meshes.append(copick_mesh)
+                        total_vertices += len(surface_mesh.vertices)
+                        total_faces += len(surface_mesh.faces)
+                        logger.info(
+                            f"Created individual surface mesh {i} with {len(surface_mesh.vertices)} vertices",
+                        )
+                    except Exception as e:
+                        logger.error(f"Failed to create mesh {i}: {e}")
+                        continue
+
+                # Return the first mesh and total stats
+                if created_meshes:
+                    stats = {"vertices_created": total_vertices, "faces_created": total_faces}
+                    return created_meshes[0], stats
+                else:
+                    return None
+            else:
+                # Create meshes from all clusters and combine them
+                all_meshes = []
+                for cluster_points in point_clusters:
+                    surface_mesh = create_surface_from_points(cluster_points)
+                    all_meshes.append(surface_mesh)
+
+                # Combine all meshes
+                combined_mesh = tm.util.concatenate(all_meshes)
+        else:
+            # Use largest cluster
+            cluster_sizes = [len(cluster) for cluster in point_clusters]
+            largest_cluster_idx = np.argmax(cluster_sizes)
+            points_to_use = point_clusters[largest_cluster_idx]
+            logger.info(f"Using largest cluster with {len(points_to_use)} points")
+
+            combined_mesh = create_surface_from_points(points_to_use)
+    else:
+        # Use all points without clustering
+        combined_mesh = create_surface_from_points(points)
+
+    # Store mesh and return stats
+    try:
+        return store_mesh_with_stats(run, combined_mesh, object_name, session_id, user_id, "surface")
+    except Exception as e:
+        logger.critical(f"Error creating mesh: {e}")
+        return None
+
+
+# Create worker function using common infrastructure
+_surface_from_picks_worker = create_batch_worker(surface_from_picks, "surface", min_points=3)
+
+
+# Create batch converter using common infrastructure
+surface_from_picks_batch = create_batch_converter(
+    surface_from_picks,
+    "Converting picks to surface meshes",
+    "surface",
+    min_points=3,
+)
+
+# Lazy batch converter for new architecture
+surface_from_picks_lazy_batch = create_lazy_batch_converter(
+    converter_func=surface_from_picks,
+    task_description="Converting picks to surface meshes",
+)

copick_utils/features/skimage.py

@@ -1,19 +1,30 @@
 import numpy as np
-from skimage.feature import multiscale_basic_features
 import zarr
 from numcodecs import Blosc
+from skimage.feature import multiscale_basic_features
+
 
-def compute_skimage_features(tomogram, feature_type, copick_root, intensity=True, edges=True, texture=True, sigma_min=0.5, sigma_max=16.0, feature_chunk_size=None):
+def compute_skimage_features(
+    tomogram,
+    feature_type,
+    copick_root,
+    intensity=True,
+    edges=True,
+    texture=True,
+    sigma_min=0.5,
+    sigma_max=16.0,
+    feature_chunk_size=None,
+):
     """
     Processes the tomogram chunkwise and computes the multiscale basic features.
     Allows for optional feature chunk size.
     """
-    image = zarr.open(tomogram.zarr(), mode='r')['0']
+    image = zarr.open(tomogram.zarr(), mode="r")["0"]
     input_chunk_size = feature_chunk_size if feature_chunk_size else image.chunks
     chunk_size = input_chunk_size if len(input_chunk_size) == 3 else input_chunk_size[1:]
-    
+
     overlap = int(chunk_size[0] / 2)
-    
+
     print(f"Processing image with shape {image.shape}")
     print(f"Using chunk size: {chunk_size}, overlap: {overlap}")
 
@@ -25,7 +36,7 @@ def compute_skimage_features(tomogram, feature_type, copick_root, intensity=True
         edges=edges,
         texture=texture,
         sigma_min=sigma_min,
-        sigma_max=sigma_max
+        sigma_max=sigma_max,
     )
     num_features = test_features.shape[-1]
 
@@ -43,10 +54,10 @@ def compute_skimage_features(tomogram, feature_type, copick_root, intensity=True
     out_array = zarr.create(
         shape=(num_features, *image.shape),
         chunks=feature_chunk_size,
-        dtype='float32',
-        compressor=Blosc(cname='zstd', clevel=3, shuffle=2),
+        dtype="float32",
+        compressor=Blosc(cname="zstd", clevel=3, shuffle=2),
         store=feature_store,
-        overwrite=True
+        overwrite=True,
     )
 
     # Process each chunk
@@ -67,7 +78,7 @@ def compute_skimage_features(tomogram, feature_type, copick_root, intensity=True
                     edges=edges,
                     texture=texture,
                     sigma_min=sigma_min,
-                    sigma_max=sigma_max
+                    sigma_max=sigma_max,
                 )
 
                 # Adjust indices for overlap
@@ -78,7 +89,12 @@ def compute_skimage_features(tomogram, feature_type, copick_root, intensity=True
                 # Ensure contiguous array and correct slicing
                 contiguous_chunk = np.ascontiguousarray(chunk_features[z_slice, y_slice, x_slice].transpose(3, 0, 1, 2))
 
-                out_array[0:num_features, z:z + chunk_size[0], y:y + chunk_size[1], x:x + chunk_size[2]] = contiguous_chunk
+                out_array[
+                    0:num_features,
+                    z : z + chunk_size[0],
+                    y : y + chunk_size[1],
+                    x : x + chunk_size[2],
+                ] = contiguous_chunk
 
     print(f"Features saved under feature type '{feature_type}'")
     return copick_features
@@ -91,6 +107,10 @@ if __name__ == "__main__":
         tomogram=tomo,
         feature_type="skimageFeatures",
         copick_root=root,
-        intensity=True, edges=True, texture=True, sigma_min=0.5, sigma_max=16.0,
-        feature_chunk_size=None  # Default to detected chunk size
+        intensity=True,
+        edges=True,
+        texture=True,
+        sigma_min=0.5,
+        sigma_max=16.0,
+        feature_chunk_size=None,  # Default to detected chunk size
     )

copick_utils/io/readers.py

@@ -1,33 +1,34 @@
 import numpy as np
 
-def tomogram(run, 
-             voxel_size: float = 10, 
-             algorithm: str = 'wbp',
-             raise_error: bool = False):
-    
+
+def tomogram(run, voxel_size: float = 10, algorithm: str = "wbp", raise_error: bool = False):
     voxel_spacing_obj = run.get_voxel_spacing(voxel_size)
 
     if voxel_spacing_obj is None:
         # Query Avaiable Voxel Spacings
         availableVoxelSpacings = [tomo.voxel_size for tomo in run.voxel_spacings]
 
-        # Report to the user which voxel spacings they can use 
-        message = (f"[Warning] No tomogram found for {run.name} with voxel size {voxel_size} and tomogram type {algorithm}"
-                   f"Available spacings are: {', '.join(map(str, availableVoxelSpacings))}" ) 
+        # Report to the user which voxel spacings they can use
+        message = (
+            f"[Warning] No tomogram found for {run.name} with voxel size {voxel_size} and tomogram type {algorithm}"
+            f"Available spacings are: {', '.join(map(str, availableVoxelSpacings))}"
+        )
         if raise_error:
             raise ValueError(message)
         else:
             print(message)
             return None
-    
+
     tomogram = voxel_spacing_obj.get_tomogram(algorithm)
     if tomogram is None:
         # Get available algorithms
         availableAlgorithms = [tomo.tomo_type for tomo in run.get_voxel_spacing(voxel_size).tomograms]
-        
+
         # Report to the user which algorithms are available
-        message = (f"[Warning] No tomogram found for {run.name} with voxel size {voxel_size} and tomogram type {algorithm}"
-                   f"Available algorithms are: {', '.join(availableAlgorithms)}")
+        message = (
+            f"[Warning] No tomogram found for {run.name} with voxel size {voxel_size} and tomogram type {algorithm}"
+            f"Available algorithms are: {', '.join(availableAlgorithms)}"
+        )
         if raise_error:
             raise ValueError(message)
         else:
@@ -36,32 +37,29 @@ def tomogram(run,
 
     return tomogram.numpy()
 
-def segmentation(run, 
-                 voxel_spacing: float,
-                 segmentation_name: str, 
-                 session_id=None,
-                 user_id=None,
-                 raise_error = False):
 
-    seg = run.get_segmentations(name=segmentation_name, 
-                                session_id = session_id,
-                                user_id = user_id,
-                                voxel_size = voxel_spacing)
+def segmentation(run, voxel_spacing: float, segmentation_name: str, session_id=None, user_id=None, raise_error=False):
+    seg = run.get_segmentations(
+        name=segmentation_name,
+        session_id=session_id,
+        user_id=user_id,
+        voxel_size=voxel_spacing,
+    )
 
     # No Segmentations Are Available, Result in Error
     if len(seg) == 0:
         # Get all available segmentations with their metadata
         available_segs = run.get_segmentations(voxel_size=voxel_spacing)
         seg_info = [(s.name, s.user_id, s.session_id) for s in available_segs]
-        
+
         # Format the information for display
-        seg_details = [f"(name: {name}, user_id: {uid}, session_id: {sid})" 
-                      for name, uid, sid in seg_info]
-        
-        message = ( f'\nNo segmentation found matching:\n'
-                    f'  name: {segmentation_name}, user_id: {user_id}, session_id: {session_id}\n'
-                    f'Available segmentations in {run.name} are:\n  ' + 
-                    '\n  '.join(seg_details) )
+        seg_details = [f"(name: {name}, user_id: {uid}, session_id: {sid})" for name, uid, sid in seg_info]
+
+        message = (
+            f"\nNo segmentation found matching:\n"
+            f"  name: {segmentation_name}, user_id: {user_id}, session_id: {session_id}\n"
+            f"Available segmentations in {run.name} are:\n  " + "\n  ".join(seg_details)
+        )
         if raise_error:
             raise ValueError(message)
         else:
@@ -70,20 +68,24 @@ def segmentation(run,
 
     # No Segmentations Are Available, Result in Error
     if len(seg) > 1:
-        print(f'[Warning] More Than 1 Segmentation is Available for the Query Information. '
-              f'Available Segmentations are: {seg} '
-              f'Defaulting to Loading: {seg[0]}\n')
+        print(
+            f"[Warning] More Than 1 Segmentation is Available for the Query Information. "
+            f"Available Segmentations are: {seg} "
+            f"Defaulting to Loading: {seg[0]}\n",
+        )
     seg = seg[0]
 
     return seg.numpy()
 
-def coordinates(run,                     # CoPick run object containing the segmentation data
-                name: str,               # Name of the object or protein for which coordinates are being extracted
-                user_id: str,            # Identifier of the user that generated the picks
-                session_id: str = None,  # Identifier of the session that generated the picks
-                voxel_size: float = 10,  # Voxel size of the tomogram, used for scaling the coordinates
-                raise_error: bool = False):
-                           
+
+def coordinates(
+    run,  # CoPick run object containing the segmentation data
+    name: str,  # Name of the object or protein for which coordinates are being extracted
+    user_id: str,  # Identifier of the user that generated the picks
+    session_id: str = None,  # Identifier of the session that generated the picks
+    voxel_size: float = 10,  # Voxel size of the tomogram, used for scaling the coordinates
+    raise_error: bool = False,
+):
     # Retrieve the pick points associated with the specified object and user ID
     picks = run.get_picks(object_name=name, user_id=user_id, session_id=session_id)
 
@@ -92,15 +94,15 @@ def coordinates(run,                     # CoPick run object containing the segm
 
         available_picks = run.get_picks()
         picks_info = [(s.pickable_object_name, s.user_id, s.session_id) for s in available_picks]
-        
+
         # Format the information for display
-        picks_details = [f"(name: {name}, user_id: {uid}, session_id: {sid})" 
-                      for name, uid, sid in picks_info]
-        
-        message = ( f'\nNo picks found matching:\n'
-                    f'  name: {name}, user_id: {user_id}, session_id: {session_id}\n'
-                    f'Available picks are:\n  '  
-                    + '\n  '.join(picks_details) )
+        picks_details = [f"(name: {name}, user_id: {uid}, session_id: {sid})" for name, uid, sid in picks_info]
+
+        message = (
+            f"\nNo picks found matching:\n"
+            f"  name: {name}, user_id: {user_id}, session_id: {session_id}\n"
+            f"Available picks are:\n  " + "\n  ".join(picks_details)
+        )
         if raise_error:
             raise ValueError(message)
         else:
@@ -109,24 +111,25 @@ def coordinates(run,                     # CoPick run object containing the segm
     elif len(picks) > 1:
         # Format pick information for display
         picks_info = [(p.pickable_object_name, p.user_id, p.session_id) for p in picks]
-        picks_details = [f"(name: {name}, user_id: {uid}, session_id: {sid})" 
-                        for name, uid, sid in picks_info]
+        picks_details = [f"(name: {name}, user_id: {uid}, session_id: {sid})" for name, uid, sid in picks_info]
 
-        print(f'[Warning] More than 1 pick is available for the query information.'
-              f'\nAvailable picks are:\n  ' + 
-              '\n  '.join(picks_details) +
-              f'\nDefaulting to loading:\n {picks[0]}\n')
+        print(
+            "[Warning] More than 1 pick is available for the query information."
+            "\nAvailable picks are:\n  " + "\n  ".join(picks_details) + f"\nDefaulting 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
+    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
+        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