diff --git a/src/copick_utils/cli/clipmesh.py b/src/copick_utils/cli/clipmesh.py index 980f401..521ed1b 100644 --- a/src/copick_utils/cli/clipmesh.py +++ b/src/copick_utils/cli/clipmesh.py @@ -8,9 +8,11 @@ from copick_utils.cli.util import ( add_distance_options, add_input_option, + add_invert_option, add_output_option, add_reference_mesh_option, add_reference_seg_option, + add_reference_tomogram_option, add_workers_option, ) from copick_utils.util.config_models import create_reference_config @@ -30,11 +32,13 @@ help="Specific run names to process (default: all runs).", ) @add_input_option("mesh") -@optgroup.group("\nReference Options", help="Options for reference surface (provide either mesh or segmentation).") +@optgroup.group("\nReference Options", help="Options for reference surface (provide mesh, segmentation, or tomogram).") @add_reference_mesh_option(required=False) @add_reference_seg_option(required=False) +@add_reference_tomogram_option(required=False) @optgroup.group("\nTool Options", help="Options related to this tool.") @add_distance_options +@add_invert_option @add_workers_option @optgroup.group("\nOutput Options", help="Options related to output meshes.") @add_output_option("mesh", default_tool="clipmesh") @@ -52,8 +56,10 @@ def clipmesh( input_uri, ref_mesh_uri, ref_seg_uri, + ref_tomo_uri, max_distance, mesh_voxel_spacing, + invert, workers, output_uri, individual_meshes, @@ -66,35 +72,50 @@ def clipmesh( URI Format: Meshes: object_name:user_id/session_id Segmentations: name:user_id/session_id@voxel_spacing + Tomograms: tomo_type@voxel_spacing \b - The reference surface can be either a mesh or a segmentation. - Only mesh vertices within the specified distance will be kept. + The reference surface can be a mesh, segmentation, or tomogram boundary. + By default, vertices within the specified distance will be kept. + Use --invert to keep vertices beyond the specified distance instead. \b Examples: - # Limit mesh to vertices near reference mesh surface - copick logical clipmesh -i "membrane:user1/full-001" -rm "boundary:user1/boundary-001" -o "membrane:clipmesh/limited-001" --max-distance 50.0 + # Keep vertices within 50Å of mesh reference + copick logical clipmesh -i "membrane:user1/full-001" -rm "boundary:user1/boundary-001" -o "membrane:clipmesh/near-001" --max-distance 50.0 - # Limit using segmentation as reference - copick logical clipmesh -i "membrane:user1/full-001" -rs "mask:user1/mask-001@10.0" -o "membrane:clipmesh/limited-001" --max-distance 100.0 + # Remove vertices within 50Å of tomogram boundary (keep interior vertices) + copick logical clipmesh -i "membrane:user1/full-001" -rt "wbp@10.0" -o "membrane:clipmesh/interior-001" --max-distance 50.0 --invert + + # Keep vertices beyond 100Å from segmentation reference + copick logical clipmesh -i "membrane:user1/full-001" -rs "mask:user1/mask-001@10.0" -o "membrane:clipmesh/far-001" --max-distance 100.0 --invert """ from copick_utils.logical.distance_operations import limit_mesh_by_distance_lazy_batch logger = get_logger(__name__, debug=debug) # Validate that exactly one reference type is provided - if not ref_mesh_uri and not ref_seg_uri: - raise click.BadParameter("Must provide either --ref-mesh or --ref-seg") - if ref_mesh_uri and ref_seg_uri: - raise click.BadParameter("Cannot provide both --ref-mesh and --ref-seg") + ref_count = sum(1 for ref in [ref_mesh_uri, ref_seg_uri, ref_tomo_uri] if ref) + if ref_count == 0: + raise click.BadParameter("Must provide one of --ref-mesh, --ref-seg, or --ref-tomogram") + if ref_count > 1: + raise click.BadParameter( + "Cannot provide multiple reference options. Choose one of --ref-mesh, --ref-seg, or --ref-tomogram", + ) root = copick.from_file(config) run_names_list = list(run_names) if run_names else None # Determine reference type and URI - reference_uri = ref_mesh_uri or ref_seg_uri - reference_type = "mesh" if ref_mesh_uri else "segmentation" + if ref_mesh_uri: + reference_uri = ref_mesh_uri + reference_type = "mesh" + elif ref_seg_uri: + reference_uri = ref_seg_uri + reference_type = "segmentation" + else: + reference_uri = ref_tomo_uri + reference_type = "tomogram" # Create config directly from URIs with additional params for distance operations try: @@ -105,7 +126,11 @@ def clipmesh( output_type="mesh", reference_uri=reference_uri, reference_type=reference_type, - additional_params={"max_distance": max_distance, "mesh_voxel_spacing": mesh_voxel_spacing}, + additional_params={ + "max_distance": max_distance, + "mesh_voxel_spacing": mesh_voxel_spacing, + "invert": invert, + }, command_name="clipmesh", ) except ValueError as e: @@ -120,11 +145,14 @@ def clipmesh( logger.info(f"Source mesh pattern: {input_params['user_id']}/{input_params['session_id']}") if reference_type == "mesh": logger.info(f"Reference mesh: {ref_params['object_name']} ({ref_params['user_id']}/{ref_params['session_id']})") - else: + elif reference_type == "segmentation": logger.info( f"Reference segmentation: {ref_params['name']} ({ref_params['user_id']}/{ref_params['session_id']})", ) + else: + logger.info(f"Reference tomogram boundary: {ref_params['tomo_type']}@{ref_params['voxel_spacing']}") logger.info(f"Maximum distance: {max_distance} angstroms") + logger.info(f"Mode: {'keep beyond distance (invert)' if invert else 'keep within distance'}") logger.info( f"Target mesh template: {output_params['object_name']} ({output_params['user_id']}/{output_params['session_id']})", ) diff --git a/src/copick_utils/cli/clippicks.py b/src/copick_utils/cli/clippicks.py index 9c0e23c..fe1c241 100644 --- a/src/copick_utils/cli/clippicks.py +++ b/src/copick_utils/cli/clippicks.py @@ -10,9 +10,11 @@ from copick_utils.cli.util import ( add_distance_options, add_input_option, + add_invert_option, add_output_option, add_reference_mesh_option, add_reference_seg_option, + add_reference_tomogram_option, add_workers_option, ) from copick_utils.util.config_models import create_reference_config @@ -32,11 +34,13 @@ help="Specific run names to process (default: all runs).", ) @add_input_option("picks") -@optgroup.group("\nReference Options", help="Options for reference surface (provide either mesh or segmentation).") +@optgroup.group("\nReference Options", help="Options for reference surface (provide mesh, segmentation, or tomogram).") @add_reference_mesh_option(required=False) @add_reference_seg_option(required=False) +@add_reference_tomogram_option(required=False) @optgroup.group("\nTool Options", help="Options related to this tool.") @add_distance_options +@add_invert_option @add_workers_option @optgroup.group("\nOutput Options", help="Options related to output picks.") @add_output_option("picks", default_tool="clippicks") @@ -47,8 +51,10 @@ def clippicks( input_uri, ref_mesh_uri, ref_seg_uri, + ref_tomo_uri, max_distance, mesh_voxel_spacing, + invert, workers, output_uri, debug, @@ -61,35 +67,50 @@ def clippicks( Picks: object_name:user_id/session_id Meshes: object_name:user_id/session_id Segmentations: name:user_id/session_id@voxel_spacing + Tomograms: tomo_type@voxel_spacing \b - The reference surface can be either a mesh or a segmentation. - Only picks within the specified distance will be kept. + The reference surface can be a mesh, segmentation, or tomogram boundary. + By default, picks within the specified distance will be kept. + Use --invert to keep picks beyond the specified distance instead. \b Examples: - # Limit picks to those near reference mesh surface - copick logical clippicks -i "ribosome:user1/all-001" -rm "boundary:user1/boundary-001" -o "ribosome:clippicks/limited-001" --max-distance 50.0 + # Keep picks within 50Å of mesh reference + copick logical clippicks -i "ribosome:user1/all-001" -rm "boundary:user1/boundary-001" -o "ribosome:clippicks/near-001" --max-distance 50.0 - # Limit using segmentation as reference - copick logical clippicks -i "ribosome:user1/all-001" -rs "mask:user1/mask-001@10.0" -o "ribosome:clippicks/limited-001" --max-distance 100.0 + # Remove picks within 50Å of tomogram boundary (keep interior picks) + copick logical clippicks -i "ribosome:user1/all-001" -rt "wbp@10.0" -o "ribosome:clippicks/interior-001" --max-distance 50.0 --invert + + # Keep picks beyond 100Å from segmentation reference + copick logical clippicks -i "ribosome:user1/all-001" -rs "mask:user1/mask-001@10.0" -o "ribosome:clippicks/far-001" --max-distance 100.0 --invert """ from copick_utils.logical.distance_operations import limit_picks_by_distance_lazy_batch logger = get_logger(__name__, debug=debug) # Validate that exactly one reference type is provided - if not ref_mesh_uri and not ref_seg_uri: - raise click.BadParameter("Must provide either --ref-mesh or --ref-seg") - if ref_mesh_uri and ref_seg_uri: - raise click.BadParameter("Cannot provide both --ref-mesh and --ref-seg") + ref_count = sum(1 for ref in [ref_mesh_uri, ref_seg_uri, ref_tomo_uri] if ref) + if ref_count == 0: + raise click.BadParameter("Must provide one of --ref-mesh, --ref-seg, or --ref-tomogram") + if ref_count > 1: + raise click.BadParameter( + "Cannot provide multiple reference options. Choose one of --ref-mesh, --ref-seg, or --ref-tomogram", + ) root = copick.from_file(config) run_names_list = list(run_names) if run_names else None # Determine reference type and URI - reference_uri = ref_mesh_uri or ref_seg_uri - reference_type = "mesh" if ref_mesh_uri else "segmentation" + if ref_mesh_uri: + reference_uri = ref_mesh_uri + reference_type = "mesh" + elif ref_seg_uri: + reference_uri = ref_seg_uri + reference_type = "segmentation" + else: + reference_uri = ref_tomo_uri + reference_type = "tomogram" # Create config directly from URIs with additional params try: @@ -100,7 +121,11 @@ def clippicks( output_type="picks", reference_uri=reference_uri, reference_type=reference_type, - additional_params={"max_distance": max_distance, "mesh_voxel_spacing": mesh_voxel_spacing}, + additional_params={ + "max_distance": max_distance, + "mesh_voxel_spacing": mesh_voxel_spacing, + "invert": invert, + }, command_name="clippicks", ) except ValueError as e: @@ -115,11 +140,14 @@ def clippicks( logger.info(f"Source picks pattern: {input_params['user_id']}/{input_params['session_id']}") if reference_type == "mesh": logger.info(f"Reference mesh: {ref_params['object_name']} ({ref_params['user_id']}/{ref_params['session_id']})") - else: + elif reference_type == "segmentation": logger.info( f"Reference segmentation: {ref_params['name']} ({ref_params['user_id']}/{ref_params['session_id']})", ) + else: + logger.info(f"Reference tomogram boundary: {ref_params['tomo_type']}@{ref_params['voxel_spacing']}") logger.info(f"Maximum distance: {max_distance} angstroms") + logger.info(f"Mode: {'keep beyond distance (invert)' if invert else 'keep within distance'}") logger.info( f"Target picks template: {output_params['object_name']} ({output_params['user_id']}/{output_params['session_id']})", ) diff --git a/src/copick_utils/cli/clipseg.py b/src/copick_utils/cli/clipseg.py index e594139..b7ffa9d 100644 --- a/src/copick_utils/cli/clipseg.py +++ b/src/copick_utils/cli/clipseg.py @@ -8,9 +8,11 @@ from copick_utils.cli.util import ( add_distance_options, add_input_option, + add_invert_option, add_output_option, add_reference_mesh_option, add_reference_seg_option, + add_reference_tomogram_option, add_workers_option, ) from copick_utils.util.config_models import create_reference_config @@ -30,11 +32,13 @@ help="Specific run names to process (default: all runs).", ) @add_input_option("segmentation") -@optgroup.group("\nReference Options", help="Options for reference surface (provide either mesh or segmentation).") +@optgroup.group("\nReference Options", help="Options for reference surface (provide mesh, segmentation, or tomogram).") @add_reference_mesh_option(required=False) @add_reference_seg_option(required=False) +@add_reference_tomogram_option(required=False) @optgroup.group("\nTool Options", help="Options related to this tool.") @add_distance_options +@add_invert_option @add_workers_option @optgroup.group("\nOutput Options", help="Options related to output segmentations.") @add_output_option("segmentation", default_tool="clipseg") @@ -45,8 +49,10 @@ def clipseg( input_uri, ref_mesh_uri, ref_seg_uri, + ref_tomo_uri, max_distance, mesh_voxel_spacing, + invert, workers, output_uri, debug, @@ -58,35 +64,50 @@ def clipseg( URI Format: Segmentations: name:user_id/session_id@voxel_spacing Meshes: object_name:user_id/session_id + Tomograms: tomo_type@voxel_spacing \b - The reference surface can be either a mesh or another segmentation. - Only segmentation voxels within the specified distance will be kept. + The reference surface can be a mesh, segmentation, or tomogram boundary. + By default, voxels within the specified distance will be kept. + Use --invert to keep voxels beyond the specified distance instead. \b Examples: - # Limit segmentation to voxels near reference mesh - copick logical clipseg -i "membrane:user1/full-001@10.0" -rm "boundary:user1/boundary-001" -o "membrane:clipseg/limited-001@10.0" --max-distance 50.0 + # Keep voxels within 50Å of mesh reference + copick logical clipseg -i "membrane:user1/full-001@10.0" -rm "boundary:user1/boundary-001" -o "membrane:clipseg/near-001@10.0" --max-distance 50.0 - # Limit using another segmentation as reference - copick logical clipseg -i "membrane:user1/full-001@10.0" -rs "mask:user1/mask-001@10.0" -o "membrane:clipseg/limited-001@10.0" --max-distance 100.0 + # Remove voxels within 50Å of tomogram boundary (keep interior voxels) + copick logical clipseg -i "membrane:user1/full-001@10.0" -rt "wbp@10.0" -o "membrane:clipseg/interior-001@10.0" --max-distance 50.0 --invert + + # Keep voxels beyond 100Å from segmentation reference + copick logical clipseg -i "membrane:user1/full-001@10.0" -rs "mask:user1/mask-001@10.0" -o "membrane:clipseg/far-001@10.0" --max-distance 100.0 --invert """ from copick_utils.logical.distance_operations import limit_segmentation_by_distance_lazy_batch logger = get_logger(__name__, debug=debug) # Validate that exactly one reference type is provided - if not ref_mesh_uri and not ref_seg_uri: - raise click.BadParameter("Must provide either --ref-mesh or --ref-seg") - if ref_mesh_uri and ref_seg_uri: - raise click.BadParameter("Cannot provide both --ref-mesh and --ref-seg") + ref_count = sum(1 for ref in [ref_mesh_uri, ref_seg_uri, ref_tomo_uri] if ref) + if ref_count == 0: + raise click.BadParameter("Must provide one of --ref-mesh, --ref-seg, or --ref-tomogram") + if ref_count > 1: + raise click.BadParameter( + "Cannot provide multiple reference options. Choose one of --ref-mesh, --ref-seg, or --ref-tomogram", + ) root = copick.from_file(config) run_names_list = list(run_names) if run_names else None # Determine reference type and URI - reference_uri = ref_mesh_uri or ref_seg_uri - reference_type = "mesh" if ref_mesh_uri else "segmentation" + if ref_mesh_uri: + reference_uri = ref_mesh_uri + reference_type = "mesh" + elif ref_seg_uri: + reference_uri = ref_seg_uri + reference_type = "segmentation" + else: + reference_uri = ref_tomo_uri + reference_type = "tomogram" # Extract voxel_spacing and multilabel from output for additional_params output_params_temp = parse_copick_uri(output_uri, "segmentation") @@ -109,6 +130,7 @@ def clipseg( "mesh_voxel_spacing": mesh_voxel_spacing, "voxel_spacing": voxel_spacing_output, "is_multilabel": multilabel_output, + "invert": invert, }, command_name="clipseg", ) @@ -124,11 +146,14 @@ def clipseg( logger.info(f"Source segmentation pattern: {input_params['user_id']}/{input_params['session_id']}") if reference_type == "mesh": logger.info(f"Reference mesh: {ref_params['object_name']} ({ref_params['user_id']}/{ref_params['session_id']})") - else: + elif reference_type == "segmentation": logger.info( f"Reference segmentation: {ref_params['name']} ({ref_params['user_id']}/{ref_params['session_id']})", ) + else: + logger.info(f"Reference tomogram boundary: {ref_params['tomo_type']}@{ref_params['voxel_spacing']}") logger.info(f"Maximum distance: {max_distance} angstroms") + logger.info(f"Mode: {'keep beyond distance (invert)' if invert else 'keep within distance'}") logger.info( f"Target segmentation template: {output_params['name']} ({output_params['user_id']}/{output_params['session_id']})", ) diff --git a/src/copick_utils/cli/util.py b/src/copick_utils/cli/util.py index 1ffdeaf..efe0736 100644 --- a/src/copick_utils/cli/util.py +++ b/src/copick_utils/cli/util.py @@ -578,3 +578,72 @@ def add_tomogram_option_decorator(_func: click.Command) -> click.Command: return add_tomogram_option_decorator else: return add_tomogram_option_decorator(func) + + +def add_reference_tomogram_option(func: click.Command = None, required: bool = False) -> Callable: + """ + Add --ref-tomogram/-rt option for using tomogram boundaries as a reference surface. + + When specified, the tomogram volume boundaries (6 faces of the bounding box) are used + as the reference surface for distance calculations. + + Tomogram URI format: tomo_type@voxel_spacing + Example: "wbp@10.0" + + Args: + func (click.Command, optional): The Click command to which the option will be added. + required (bool): Whether the option is required. Default is False. + + Returns: + Callable: The Click command with the reference tomogram option added. + """ + + def add_reference_tomogram_option_decorator(_func: click.Command) -> click.Command: + """Add --ref-tomogram option to command.""" + opt = optgroup.option( + "--ref-tomogram", + "-rt", + "ref_tomo_uri", + required=required, + help="Reference tomogram boundary URI (format: tomo_type@voxel_spacing). " + "Uses tomogram volume boundaries as reference surface. Example: 'wbp@10.0'", + ) + return opt(_func) + + if func is None: + return add_reference_tomogram_option_decorator + else: + return add_reference_tomogram_option_decorator(func) + + +def add_invert_option(func: click.Command = None) -> Callable: + """ + Add --invert flag to invert distance filtering behavior. + + When False (default): Keep data WITHIN max_distance of reference surface. + When True: Keep data BEYOND max_distance of reference surface (remove data within distance). + + This flag applies uniformly to all reference types (mesh, segmentation, tomogram). + + Args: + func (click.Command, optional): The Click command to which the option will be added. + + Returns: + Callable: The Click command with the invert option added. + """ + + def add_invert_option_decorator(_func: click.Command) -> click.Command: + """Add --invert option to command.""" + opt = optgroup.option( + "--invert/--no-invert", + "invert", + is_flag=True, + default=False, + help="Invert filtering: keep data BEYOND max_distance instead of WITHIN (default: keep within).", + ) + return opt(_func) + + if func is None: + return add_invert_option_decorator + else: + return add_invert_option_decorator(func) diff --git a/src/copick_utils/converters/lazy_converter.py b/src/copick_utils/converters/lazy_converter.py index 02e091d..5cc5507 100644 --- a/src/copick_utils/converters/lazy_converter.py +++ b/src/copick_utils/converters/lazy_converter.py @@ -205,7 +205,25 @@ def add_references_to_tasks( """ reference_type = reference_config.reference_type - # Find reference objects + # Handle tomogram reference type (uses tomogram boundaries, not an object) + if reference_type == "tomogram": + # For tomogram references, we pass tomo_type and voxel_spacing as reference_tomogram_info + # The actual tomogram bounds are computed in the converter function + augmented_tasks = [] + for task in tasks: + task["reference_mesh"] = None + task["reference_segmentation"] = None + task["reference_tomogram_info"] = (reference_config.tomo_type, reference_config.voxel_spacing) + + # Add additional reference parameters + for key, value in reference_config.additional_params.items(): + task[key] = value + + augmented_tasks.append(task) + + return augmented_tasks + + # Find reference objects for mesh/segmentation if reference_type == "mesh": ref_objects = run.get_meshes( object_name=reference_config.object_name, @@ -233,6 +251,7 @@ def add_references_to_tasks( for task in tasks: task[ref_key] = ref_objects[0] task[alt_key] = None + task["reference_tomogram_info"] = None # Ensure tomogram info is None for mesh/seg references # Add additional reference parameters for key, value in reference_config.additional_params.items(): diff --git a/src/copick_utils/logical/distance_operations.py b/src/copick_utils/logical/distance_operations.py index 2c2f5bf..855c502 100644 --- a/src/copick_utils/logical/distance_operations.py +++ b/src/copick_utils/logical/distance_operations.py @@ -117,6 +117,119 @@ def _create_distance_field_from_mesh( return _create_distance_field_from_segmentation(voxelized_array.astype(np.uint8), target_voxel_spacing) +def _get_tomogram_bounds( + run: "CopickRun", + tomo_type: str, + voxel_spacing: float, +) -> Tuple[Tuple[float, float, float], Tuple[float, float, float]]: + """ + Get tomogram physical bounds as ((0,0,0), (max_x, max_y, max_z)). + + Args: + run: CopickRun object + tomo_type: Type of the tomogram (e.g., "wbp") + voxel_spacing: Voxel spacing of the tomogram + + Returns: + Tuple of ((min_x, min_y, min_z), (max_x, max_y, max_z)) in physical units (angstroms) + + Raises: + ValueError: If voxel spacing or tomogram type not found + """ + import zarr + + vs = run.get_voxel_spacing(voxel_spacing) + if vs is None: + available = [v.voxel_size for v in run.voxel_spacings] + raise ValueError(f"Voxel spacing {voxel_spacing} not found in run {run.name}. Available: {available}") + + tomo = vs.get_tomogram(tomo_type) + if tomo is None: + available = [t.tomo_type for t in vs.tomograms] + raise ValueError( + f"Tomogram type '{tomo_type}' not found at voxel spacing {voxel_spacing}. Available: {available}", + ) + + # Get shape from zarr (z, y, x order) + zarr_array = zarr.open(tomo.zarr())["0"] + shape_zyx = zarr_array.shape + shape_xyz = shape_zyx[::-1] # Convert to (x, y, z) + + # Physical bounds in angstroms + max_bound = tuple(np.array(shape_xyz) * voxel_spacing) + return ((0.0, 0.0, 0.0), max_bound) + + +def _compute_distances_to_tomogram_boundary( + positions: np.ndarray, + tomogram_bounds: Tuple[Tuple[float, float, float], Tuple[float, float, float]], +) -> np.ndarray: + """ + Compute distance from each position to nearest tomogram boundary. + + The distance is the minimum distance to any of the 6 faces of the tomogram bounding box. + + Args: + positions: Array of shape (N, 3) with positions in physical coordinates (angstroms) + tomogram_bounds: Tuple of ((min_x, min_y, min_z), (max_x, max_y, max_z)) + + Returns: + Array of shape (N,) with distance to nearest boundary for each position + """ + min_bound = np.array(tomogram_bounds[0]) + max_bound = np.array(tomogram_bounds[1]) + + # Distance to each of the 6 boundaries + dist_to_min = positions - min_bound # Distance to x=0, y=0, z=0 planes + dist_to_max = max_bound - positions # Distance to x=max, y=max, z=max planes + + # Minimum distance across all 6 boundaries + all_distances = np.column_stack([dist_to_min, dist_to_max]) + return np.min(all_distances, axis=1) + + +def _create_distance_field_from_tomogram_boundary( + tomogram_shape: Tuple[int, int, int], + voxel_spacing: float, +) -> np.ndarray: + """ + Create distance field where each voxel contains the distance to the nearest tomogram boundary. + + This is used for segmentation and mesh filtering where a full distance field is needed. + + Args: + tomogram_shape: Shape of the tomogram (z, y, x) or (x, y, z) in voxels + voxel_spacing: Voxel spacing in physical units + + Returns: + Distance field array with distances in physical units + """ + # Create coordinate grids + z, y, x = np.indices(tomogram_shape) + + # Distance to each boundary (in voxels) + dist_to_x_min = x + dist_to_x_max = tomogram_shape[2] - 1 - x + dist_to_y_min = y + dist_to_y_max = tomogram_shape[1] - 1 - y + dist_to_z_min = z + dist_to_z_max = tomogram_shape[0] - 1 - z + + # Minimum distance to any boundary + distance_field_voxels = np.minimum.reduce( + [ + dist_to_x_min, + dist_to_x_max, + dist_to_y_min, + dist_to_y_max, + dist_to_z_min, + dist_to_z_max, + ], + ) + + return distance_field_voxels.astype(float) * voxel_spacing + + def limit_mesh_by_distance( mesh: "CopickMesh", run: "CopickRun", @@ -125,8 +238,10 @@ def limit_mesh_by_distance( output_user_id: str, reference_mesh: Optional["CopickMesh"] = None, reference_segmentation: Optional["CopickSegmentation"] = None, + reference_tomogram_info: Optional[Tuple[str, float]] = None, max_distance: float = 100.0, mesh_voxel_spacing: float = None, + invert: bool = False, **kwargs, ) -> Optional[Tuple["CopickMesh", Dict[str, int]]]: """ @@ -134,14 +249,16 @@ def limit_mesh_by_distance( Args: mesh: CopickMesh to limit - reference_mesh: Reference CopickMesh (either this or reference_segmentation must be provided) + reference_mesh: Reference CopickMesh (one of mesh/segmentation/tomogram must be provided) reference_segmentation: Reference CopickSegmentation + reference_tomogram_info: Tuple of (tomo_type, voxel_spacing) for tomogram boundary reference run: CopickRun object output_object_name: Name for the output mesh output_session_id: Session ID for the output mesh output_user_id: User ID for the output mesh max_distance: Maximum distance from reference surface mesh_voxel_spacing: Voxel spacing for mesh voxelization (defaults to 10.0) + invert: If False (default), keep vertices within max_distance. If True, keep vertices beyond max_distance. **kwargs: Additional keyword arguments Returns: @@ -149,8 +266,10 @@ def limit_mesh_by_distance( Stats dict contains 'vertices_created' and 'faces_created'. """ try: - if reference_mesh is None and reference_segmentation is None: - raise ValueError("Either reference_mesh or reference_segmentation must be provided") + if reference_mesh is None and reference_segmentation is None and reference_tomogram_info is None: + raise ValueError( + "One of reference_mesh, reference_segmentation, or reference_tomogram_info must be provided", + ) # Get the target mesh target_mesh = mesh.mesh @@ -165,87 +284,104 @@ def limit_mesh_by_distance( return None target_mesh = tm.util.concatenate(list(target_mesh.geometry.values())) - # Create distance field from reference - # Use mesh bounds to define coordinate space - mesh_bounds = np.array([target_mesh.vertices.min(axis=0), target_mesh.vertices.max(axis=0)]) - - # Add padding for max_distance - padding = max_distance * 1.1 - mesh_bounds[0] -= padding - mesh_bounds[1] += padding + # Handle tomogram boundary reference (direct distance calculation) + if reference_tomogram_info is not None: + tomo_type, tomo_voxel_spacing = reference_tomogram_info + tomogram_bounds = _get_tomogram_bounds(run, tomo_type, tomo_voxel_spacing) + vertex_distances = _compute_distances_to_tomogram_boundary(target_mesh.vertices, tomogram_bounds) - field_voxel_spacing = mesh_voxel_spacing if mesh_voxel_spacing is not None else 10.0 - field_size = mesh_bounds[1] - mesh_bounds[0] - field_shape = np.ceil(field_size / field_voxel_spacing).astype(int) + # Apply invert logic + final_valid = vertex_distances > max_distance if invert else vertex_distances <= max_distance - if reference_mesh is not None: - ref_mesh = reference_mesh.mesh - if ref_mesh is None: - logger.error("Could not load reference mesh data") + if not np.any(final_valid): + within_or_beyond = "beyond" if invert else "within" + logger.warning(f"No vertices {within_or_beyond} {max_distance} units of tomogram boundary") return None - if isinstance(ref_mesh, tm.Scene): - if len(ref_mesh.geometry) == 0: - logger.error("Reference mesh is empty") + # Handle mesh or segmentation reference (uses distance field) + else: + # Create distance field from reference + # Use mesh bounds to define coordinate space + mesh_bounds = np.array([target_mesh.vertices.min(axis=0), target_mesh.vertices.max(axis=0)]) + + # Add padding for max_distance + padding = max_distance * 1.1 + mesh_bounds[0] -= padding + mesh_bounds[1] += padding + + field_voxel_spacing = mesh_voxel_spacing if mesh_voxel_spacing is not None else 10.0 + field_size = mesh_bounds[1] - mesh_bounds[0] + field_shape = np.ceil(field_size / field_voxel_spacing).astype(int) + + if reference_mesh is not None: + ref_mesh = reference_mesh.mesh + if ref_mesh is None: + logger.error("Could not load reference mesh data") return None - ref_mesh = tm.util.concatenate(list(ref_mesh.geometry.values())) - # Create distance field from mesh - distance_field = _create_distance_field_from_mesh( - ref_mesh, - field_shape, - field_voxel_spacing, - mesh_voxel_spacing, - ) + if isinstance(ref_mesh, tm.Scene): + if len(ref_mesh.geometry) == 0: + logger.error("Reference mesh is empty") + return None + ref_mesh = tm.util.concatenate(list(ref_mesh.geometry.values())) + + # Create distance field from mesh + distance_field = _create_distance_field_from_mesh( + ref_mesh, + field_shape, + field_voxel_spacing, + mesh_voxel_spacing, + ) + + else: # reference_segmentation is not None + ref_seg_array = reference_segmentation.numpy() + if ref_seg_array is None or ref_seg_array.size == 0: + logger.error("Could not load reference segmentation data") + return None - else: # reference_segmentation is not None - ref_seg_array = reference_segmentation.numpy() - if ref_seg_array is None or ref_seg_array.size == 0: - logger.error("Could not load reference segmentation data") - return None + # Convert segmentation to field coordinate space + seg_indices = np.array(np.where(ref_seg_array > 0)).T + seg_physical = seg_indices * reference_segmentation.voxel_size + field_coords = np.floor((seg_physical - mesh_bounds[0]) / field_voxel_spacing).astype(int) - # Convert segmentation to field coordinate space - seg_indices = np.array(np.where(ref_seg_array > 0)).T - seg_physical = seg_indices * reference_segmentation.voxel_size - field_coords = np.floor((seg_physical - mesh_bounds[0]) / field_voxel_spacing).astype(int) - - # Create voxelized reference in field space - voxelized_ref = np.zeros(field_shape, dtype=bool) - valid_coords = (field_coords >= 0).all(axis=1) & (field_coords < field_shape).all(axis=1) - if np.any(valid_coords): - valid_field_coords = field_coords[valid_coords] - voxelized_ref[valid_field_coords[:, 0], valid_field_coords[:, 1], valid_field_coords[:, 2]] = True - - distance_field = _create_distance_field_from_segmentation( - voxelized_ref.astype(np.uint8), - field_voxel_spacing, - ) + # Create voxelized reference in field space + voxelized_ref = np.zeros(field_shape, dtype=bool) + valid_coords = (field_coords >= 0).all(axis=1) & (field_coords < field_shape).all(axis=1) + if np.any(valid_coords): + valid_field_coords = field_coords[valid_coords] + voxelized_ref[valid_field_coords[:, 0], valid_field_coords[:, 1], valid_field_coords[:, 2]] = True - # Convert mesh vertex coordinates to field indices - vertex_field_coords = (target_mesh.vertices - mesh_bounds[0]) / field_voxel_spacing - vertex_field_indices = np.floor(vertex_field_coords).astype(int) + distance_field = _create_distance_field_from_segmentation( + voxelized_ref.astype(np.uint8), + field_voxel_spacing, + ) - # Check which vertices are within field bounds - valid_vertices = (vertex_field_indices >= 0).all(axis=1) & (vertex_field_indices < field_shape).all(axis=1) + # Convert mesh vertex coordinates to field indices + vertex_field_coords = (target_mesh.vertices - mesh_bounds[0]) / field_voxel_spacing + vertex_field_indices = np.floor(vertex_field_coords).astype(int) - if not np.any(valid_vertices): - logger.warning("No mesh vertices within distance field bounds") - return None + # Check which vertices are within field bounds + valid_vertices = (vertex_field_indices >= 0).all(axis=1) & (vertex_field_indices < field_shape).all(axis=1) + + if not np.any(valid_vertices): + logger.warning("No mesh vertices within distance field bounds") + return None - # Get distances for valid vertices - valid_indices = vertex_field_indices[valid_vertices] - vertex_distances = distance_field[valid_indices[:, 0], valid_indices[:, 1], valid_indices[:, 2]] + # Get distances for valid vertices + valid_indices = vertex_field_indices[valid_vertices] + vertex_distances = distance_field[valid_indices[:, 0], valid_indices[:, 1], valid_indices[:, 2]] - # Find vertices within distance threshold - distance_valid = vertex_distances <= max_distance + # Apply invert logic: default keeps within distance, invert keeps beyond + distance_valid = vertex_distances > max_distance if invert else vertex_distances <= max_distance - # Combine bounds validity and distance validity - final_valid = np.zeros(len(target_mesh.vertices), dtype=bool) - final_valid[valid_vertices] = distance_valid + # Combine bounds validity and distance validity + final_valid = np.zeros(len(target_mesh.vertices), dtype=bool) + final_valid[valid_vertices] = distance_valid - if not np.any(final_valid): - logger.warning(f"No vertices within {max_distance} units of reference surface") - return None + if not np.any(final_valid): + within_or_beyond = "beyond" if invert else "within" + logger.warning(f"No vertices {within_or_beyond} {max_distance} units of reference surface") + return None # Create a new mesh with only valid vertices and their faces valid_vertex_indices = np.where(final_valid)[0] @@ -285,7 +421,8 @@ def limit_mesh_by_distance( shape_name="distance-limited mesh", ) - logger.info(f"Limited mesh to {stats['vertices_created']} vertices within {max_distance} units") + within_or_beyond = "beyond" if invert else "within" + logger.info(f"Limited mesh to {stats['vertices_created']} vertices {within_or_beyond} {max_distance} units") return copick_mesh, stats except Exception as e: @@ -301,10 +438,12 @@ def limit_segmentation_by_distance( output_user_id: str, reference_mesh: Optional["CopickMesh"] = None, reference_segmentation: Optional["CopickSegmentation"] = None, + reference_tomogram_info: Optional[Tuple[str, float]] = None, max_distance: float = 100.0, voxel_spacing: float = 10.0, mesh_voxel_spacing: float = None, is_multilabel: bool = False, + invert: bool = False, **kwargs, ) -> Optional[Tuple["CopickSegmentation", Dict[str, int]]]: """ @@ -312,8 +451,9 @@ def limit_segmentation_by_distance( Args: segmentation: CopickSegmentation to limit - reference_mesh: Reference CopickMesh (either this or reference_segmentation must be provided) + reference_mesh: Reference CopickMesh (one of mesh/segmentation/tomogram must be provided) reference_segmentation: Reference CopickSegmentation + reference_tomogram_info: Tuple of (tomo_type, voxel_spacing) for tomogram boundary reference run: CopickRun object output_object_name: Name for the output segmentation output_session_id: Session ID for the output segmentation @@ -322,6 +462,7 @@ def limit_segmentation_by_distance( voxel_spacing: Voxel spacing for the output segmentation mesh_voxel_spacing: Voxel spacing for mesh voxelization (defaults to target voxel spacing) is_multilabel: Whether the segmentation is multilabel + invert: If False (default), keep voxels within max_distance. If True, keep voxels beyond max_distance. **kwargs: Additional keyword arguments Returns: @@ -329,8 +470,10 @@ def limit_segmentation_by_distance( Stats dict contains 'voxels_created'. """ try: - if reference_mesh is None and reference_segmentation is None: - raise ValueError("Either reference_mesh or reference_segmentation must be provided") + if reference_mesh is None and reference_segmentation is None and reference_tomogram_info is None: + raise ValueError( + "One of reference_mesh, reference_segmentation, or reference_tomogram_info must be provided", + ) # Load target segmentation seg_array = segmentation.numpy() @@ -339,7 +482,15 @@ def limit_segmentation_by_distance( return None # Create distance field from reference - if reference_mesh is not None: + if reference_tomogram_info is not None: + # Use tomogram boundary distance field + tomo_type, tomo_voxel_spacing = reference_tomogram_info + # Verify tomogram exists (this will raise if not found) + _get_tomogram_bounds(run, tomo_type, tomo_voxel_spacing) + # Create distance field from tomogram boundary + distance_field = _create_distance_field_from_tomogram_boundary(seg_array.shape, segmentation.voxel_size) + + elif reference_mesh is not None: ref_mesh = reference_mesh.mesh if ref_mesh is None: logger.error("Could not load reference mesh data") @@ -384,14 +535,15 @@ def limit_segmentation_by_distance( # Create distance field from segmentation distance_field = _create_distance_field_from_segmentation(ref_seg_array, segmentation.voxel_size) - # Apply distance threshold to create mask - distance_mask = distance_field <= max_distance + # Apply distance threshold to create mask with invert support + distance_mask = distance_field > max_distance if invert else distance_field <= max_distance # Apply mask to target segmentation output_array = seg_array * distance_mask if np.sum(output_array > 0) == 0: - logger.warning(f"No voxels within {max_distance} units of reference surface") + within_or_beyond = "beyond" if invert else "within" + logger.warning(f"No voxels {within_or_beyond} {max_distance} units of reference surface") return None # Create output segmentation @@ -408,7 +560,8 @@ def limit_segmentation_by_distance( output_seg.from_numpy(output_array) stats = {"voxels_created": int(np.sum(output_array > 0))} - logger.info(f"Limited segmentation to {stats['voxels_created']} voxels within {max_distance} units") + within_or_beyond = "beyond" if invert else "within" + logger.info(f"Limited segmentation to {stats['voxels_created']} voxels {within_or_beyond} {max_distance} units") return output_seg, stats except Exception as e: @@ -419,13 +572,15 @@ def limit_segmentation_by_distance( def limit_picks_by_distance( picks: "CopickPicks", run: "CopickRun", - pick_object_name: str, - pick_session_id: str, - pick_user_id: str, + output_object_name: str, + output_session_id: str, + output_user_id: str, reference_mesh: Optional["CopickMesh"] = None, reference_segmentation: Optional["CopickSegmentation"] = None, + reference_tomogram_info: Optional[Tuple[str, float]] = None, max_distance: float = 100.0, mesh_voxel_spacing: float = None, + invert: bool = False, **kwargs, ) -> Optional[Tuple["CopickPicks", Dict[str, int]]]: """ @@ -433,14 +588,16 @@ def limit_picks_by_distance( Args: picks: CopickPicks to limit - reference_mesh: Reference CopickMesh (either this or reference_segmentation must be provided) + reference_mesh: Reference CopickMesh (one of mesh/segmentation/tomogram must be provided) reference_segmentation: Reference CopickSegmentation + reference_tomogram_info: Tuple of (tomo_type, voxel_spacing) for tomogram boundary reference run: CopickRun object - pick_object_name: Name for the output picks - pick_session_id: Session ID for the output picks - pick_user_id: User ID for the output picks + output_object_name: Name for the output picks + output_session_id: Session ID for the output picks + output_user_id: User ID for the output picks max_distance: Maximum distance from reference surface mesh_voxel_spacing: Voxel spacing for mesh voxelization (defaults to 10.0) + invert: If False (default), keep picks within max_distance. If True, keep picks beyond max_distance. **kwargs: Additional keyword arguments Returns: @@ -448,8 +605,10 @@ def limit_picks_by_distance( Stats dict contains 'points_created'. """ try: - if reference_mesh is None and reference_segmentation is None: - raise ValueError("Either reference_mesh or reference_segmentation must be provided") + if reference_mesh is None and reference_segmentation is None and reference_tomogram_info is None: + raise ValueError( + "One of reference_mesh, reference_segmentation, or reference_tomogram_info must be provided", + ) # Load pick data points, transforms = picks.numpy() @@ -459,92 +618,111 @@ def limit_picks_by_distance( pick_positions = points[:, :3] # Use only x, y, z coordinates - # We need a coordinate space to create the distance field - # Use the reference segmentation's coordinate space, or create one for mesh references - if reference_segmentation is not None: - ref_seg_array = reference_segmentation.numpy() - 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 - field_voxel_spacing = reference_segmentation.voxel_size - distance_field = _create_distance_field_from_segmentation(ref_seg_array, field_voxel_spacing) + # Handle tomogram boundary reference (direct distance calculation, no distance field) + if reference_tomogram_info is not None: + tomo_type, tomo_voxel_spacing = reference_tomogram_info + tomogram_bounds = _get_tomogram_bounds(run, tomo_type, tomo_voxel_spacing) + pick_distances = _compute_distances_to_tomogram_boundary(pick_positions, tomogram_bounds) - # 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) + # Apply invert logic: default keeps within distance, invert keeps beyond + final_valid = pick_distances > max_distance if invert else pick_distances <= max_distance - else: # reference_mesh is not None - ref_mesh = reference_mesh.mesh - if ref_mesh is None: - logger.error("Could not load reference mesh data") + 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 - if isinstance(ref_mesh, tm.Scene): - if len(ref_mesh.geometry) == 0: - logger.error("Reference mesh is empty") + # 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() + if ref_seg_array is None or ref_seg_array.size == 0: + logger.error("Could not load reference segmentation data") return None - ref_mesh = tm.util.concatenate(list(ref_mesh.geometry.values())) - # Define coordinate space based on mesh bounds and pick positions - all_coords = np.vstack([ref_mesh.vertices, pick_positions]) - coord_bounds = np.array([all_coords.min(axis=0), all_coords.max(axis=0)]) + # Use reference segmentation's coordinate space + field_voxel_spacing = reference_segmentation.voxel_size + distance_field = _create_distance_field_from_segmentation(ref_seg_array, field_voxel_spacing) - # Add padding for max_distance - padding = max_distance * 1.1 - coord_bounds[0] -= padding - coord_bounds[1] += padding + # 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) - field_voxel_spacing = mesh_voxel_spacing if mesh_voxel_spacing is not None else 10.0 - field_size = coord_bounds[1] - coord_bounds[0] - field_shape = np.ceil(field_size / field_voxel_spacing).astype(int) + else: # reference_mesh is not None + ref_mesh = reference_mesh.mesh + if ref_mesh is None: + logger.error("Could not load reference mesh data") + return None - # Create distance field from mesh in this coordinate space - distance_field = _create_distance_field_from_mesh( - ref_mesh, - field_shape, - field_voxel_spacing, - mesh_voxel_spacing, + if isinstance(ref_mesh, tm.Scene): + if len(ref_mesh.geometry) == 0: + logger.error("Reference mesh is empty") + return None + ref_mesh = tm.util.concatenate(list(ref_mesh.geometry.values())) + + # Define coordinate space based on mesh bounds and pick positions + all_coords = np.vstack([ref_mesh.vertices, pick_positions]) + coord_bounds = np.array([all_coords.min(axis=0), all_coords.max(axis=0)]) + + # Add padding for max_distance + padding = max_distance * 1.1 + coord_bounds[0] -= padding + coord_bounds[1] += padding + + field_voxel_spacing = mesh_voxel_spacing if mesh_voxel_spacing is not None else 10.0 + 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 + distance_field = _create_distance_field_from_mesh( + ref_mesh, + field_shape, + field_voxel_spacing, + mesh_voxel_spacing, + ) + + # Convert pick coordinates to voxel indices in this field space + pick_voxel_coords = (pick_positions - coord_bounds[0]) / field_voxel_spacing + pick_voxel_indices = np.floor(pick_voxel_coords).astype(int) + + # Check which picks are within field bounds + valid_picks = (pick_voxel_indices >= 0).all(axis=1) & (pick_voxel_indices < distance_field.shape).all( + axis=1, ) - # Convert pick coordinates to voxel indices in this field space - pick_voxel_coords = (pick_positions - coord_bounds[0]) / field_voxel_spacing - pick_voxel_indices = np.floor(pick_voxel_coords).astype(int) - - # Check which picks are within field bounds - valid_picks = (pick_voxel_indices >= 0).all(axis=1) & (pick_voxel_indices < distance_field.shape).all(axis=1) - - if not np.any(valid_picks): - logger.warning("No picks within distance field bounds") - return None + 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]] + # 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]] - # Find picks within distance threshold - distance_valid = pick_distances <= max_distance + # Apply invert logic: default keeps within distance, invert keeps beyond + distance_valid = pick_distances > max_distance if invert else pick_distances <= max_distance - # Combine bounds validity and distance validity - final_valid = np.zeros(len(points), dtype=bool) - final_valid[valid_picks] = distance_valid + # Combine bounds validity and distance validity + final_valid = np.zeros(len(points), dtype=bool) + final_valid[valid_picks] = distance_valid - if not np.any(final_valid): - logger.warning(f"No picks within {max_distance} units of reference surface") - return None + 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 # Create output picks - output_picks = run.new_picks(pick_object_name, pick_session_id, pick_user_id, exist_ok=True) + output_picks = run.new_picks(output_object_name, output_session_id, output_user_id, exist_ok=True) output_picks.from_numpy(positions=valid_points, transforms=valid_transforms) output_picks.store() stats = {"points_created": len(valid_points)} - logger.info(f"Limited picks to {stats['points_created']} points within {max_distance} units") + within_or_beyond = "beyond" if invert else "within" + logger.info(f"Limited picks to {stats['points_created']} points {within_or_beyond} {max_distance} units") return output_picks, stats except Exception as e: diff --git a/src/copick_utils/util/config_models.py b/src/copick_utils/util/config_models.py index ef0eb8a..cf4f848 100644 --- a/src/copick_utils/util/config_models.py +++ b/src/copick_utils/util/config_models.py @@ -155,42 +155,55 @@ def from_uris( class ReferenceConfig(BaseModel): """Pydantic model for reference discovery configuration.""" - reference_type: Literal["mesh", "segmentation"] + reference_type: Literal["mesh", "segmentation", "tomogram"] object_name: Optional[str] = None user_id: Optional[str] = None session_id: Optional[str] = None voxel_spacing: Optional[float] = None + tomo_type: Optional[str] = None additional_params: Dict[str, Any] = Field(default_factory=dict) @field_validator("voxel_spacing") @classmethod def validate_segmentation_voxel_spacing(cls, v, info): - """Ensure voxel_spacing is provided for segmentation references.""" + """Ensure voxel_spacing is provided for segmentation and tomogram references.""" values = info.data - if values.get("reference_type") == "segmentation" and v is None: - raise ValueError("voxel_spacing is required for segmentation references") + ref_type = values.get("reference_type") + if ref_type in ("segmentation", "tomogram") and v is None: + raise ValueError(f"voxel_spacing is required for {ref_type} references") return v @field_validator("object_name") @classmethod def validate_required_fields(cls, v, info): - """Ensure required fields are provided.""" - if v is None: + """Ensure required fields are provided for non-tomogram references.""" + values = info.data + # object_name is not required for tomogram references + if values.get("reference_type") != "tomogram" and v is None: raise ValueError("object_name is required for reference configuration") return v + @field_validator("tomo_type") + @classmethod + def validate_tomogram_fields(cls, v, info): + """Ensure tomo_type is provided for tomogram references.""" + values = info.data + if values.get("reference_type") == "tomogram" and v is None: + raise ValueError("tomo_type is required for tomogram references") + return v + @classmethod def from_uri( cls, uri: str, - reference_type: Literal["mesh", "segmentation"], + reference_type: Literal["mesh", "segmentation", "tomogram"], additional_params: Optional[Dict[str, Any]] = None, ) -> "ReferenceConfig": """Create ReferenceConfig from a URI. Args: uri: Copick URI string for the reference object. - reference_type: Type of reference ('mesh' or 'segmentation'). + reference_type: Type of reference ('mesh', 'segmentation', or 'tomogram'). additional_params: Additional parameters to include in the config. Returns: @@ -204,24 +217,40 @@ def from_uri( # Parse URI params = parse_copick_uri(uri, reference_type) - # Extract fields based on type - object_name = params.get("object_name") or params.get("name") - user_id = params["user_id"] - session_id = params["session_id"] - voxel_spacing = params.get("voxel_spacing") + if reference_type == "tomogram": + # For tomogram, extract tomo_type and voxel_spacing + tomo_type = params.get("tomo_type") + voxel_spacing = params.get("voxel_spacing") - # Convert voxel_spacing to float if it's a string - if voxel_spacing and isinstance(voxel_spacing, str) and voxel_spacing != "*": - voxel_spacing = float(voxel_spacing) + # Convert voxel_spacing to float if it's a string + if voxel_spacing and isinstance(voxel_spacing, str) and voxel_spacing != "*": + voxel_spacing = float(voxel_spacing) - return cls( - reference_type=reference_type, - object_name=object_name, - user_id=user_id, - session_id=session_id, - voxel_spacing=voxel_spacing, - additional_params=additional_params or {}, - ) + return cls( + reference_type=reference_type, + tomo_type=tomo_type, + voxel_spacing=voxel_spacing, + additional_params=additional_params or {}, + ) + else: + # For mesh/segmentation, extract standard fields + object_name = params.get("object_name") or params.get("name") + user_id = params["user_id"] + session_id = params["session_id"] + voxel_spacing = params.get("voxel_spacing") + + # Convert voxel_spacing to float if it's a string + if voxel_spacing and isinstance(voxel_spacing, str) and voxel_spacing != "*": + voxel_spacing = float(voxel_spacing) + + return cls( + reference_type=reference_type, + object_name=object_name, + user_id=user_id, + session_id=session_id, + voxel_spacing=voxel_spacing, + additional_params=additional_params or {}, + ) class TaskConfig(BaseModel): @@ -563,7 +592,7 @@ def create_reference_config( output_uri: str, output_type: Literal["picks", "mesh", "segmentation"], reference_uri: str, - reference_type: Literal["mesh", "segmentation"], + reference_type: Literal["mesh", "segmentation", "tomogram"], additional_params: Optional[Dict[str, Any]] = None, command_name: Optional[str] = None, ) -> TaskConfig: @@ -575,7 +604,7 @@ def create_reference_config( output_uri: Output copick URI string (supports smart defaults). output_type: Type of output object. reference_uri: Reference copick URI string. - reference_type: Type of reference object. + reference_type: Type of reference object ('mesh', 'segmentation', or 'tomogram'). additional_params: Additional parameters for reference config. command_name: Name of the command (used for smart defaults in output_uri). @@ -592,6 +621,17 @@ def create_reference_config( reference_type="mesh", command_name="picksin", ) + + # With tomogram boundary reference: + config = create_reference_config( + input_uri="ribosome:user1/all-001", + input_type="picks", + output_uri="ribosome", + output_type="picks", + reference_uri="wbp@10.0", + reference_type="tomogram", + command_name="clippicks", + ) """ selector_config = SelectorConfig.from_uris( input_uri=input_uri,