diff --git a/pyproject.toml b/pyproject.toml index a7c868c..835ab26 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -45,9 +45,9 @@ license = { file = "LICENSE" } keywords = ["copick", "cryoet", "cryo-et", "tomography", "annotation", "utilities"] [project.urls] -Repository = "https://github.com/KyleHarrington/copick-utils.git" -Issues = "https://github.com/KyleHarrington/copick-utils/issues" -Documentation = "https://github.com/KyleHarrington/copick-utils#readme" +Repository = "https://github.com/copick/copick-utils.git" +Issues = "https://github.com/copick/copick-utils/issues" +Documentation = "https://github.com/copick/copick-utils#readme" [project.optional-dependencies] dev = [ @@ -70,6 +70,9 @@ picks2sphere = "copick_utils.cli.conversion_commands:picks2sphere" picks2ellipsoid = "copick_utils.cli.conversion_commands:picks2ellipsoid" picks2plane = "copick_utils.cli.conversion_commands:picks2plane" picks2surface = "copick_utils.cli.conversion_commands:picks2surface" +seg2fil = "copick_utils.cli.conversion_commands:seg2fil" +fil2picks = "copick_utils.cli.conversion_commands:fil2picks" +fil2seg = "copick_utils.cli.conversion_commands:fil2seg" [project.entry-points."copick.process.commands"] separate_components = "copick_utils.cli.processing_commands:separate_components" diff --git a/src/copick_utils/cli/conversion_commands.py b/src/copick_utils/cli/conversion_commands.py index f414ee1..2a9a6be 100644 --- a/src/copick_utils/cli/conversion_commands.py +++ b/src/copick_utils/cli/conversion_commands.py @@ -3,6 +3,8 @@ This module imports all conversion commands from specialized files for better organization. """ +from copick_utils.cli.fil2picks import fil2picks +from copick_utils.cli.fil2seg import fil2seg from copick_utils.cli.mesh2caps import mesh2caps from copick_utils.cli.mesh2picks import mesh2picks from copick_utils.cli.mesh2seg import mesh2seg @@ -12,6 +14,7 @@ from copick_utils.cli.picks2seg import picks2seg from copick_utils.cli.picks2sphere import picks2sphere from copick_utils.cli.picks2surface import picks2surface +from copick_utils.cli.seg2fil import seg2fil from copick_utils.cli.seg2mesh import seg2mesh from copick_utils.cli.seg2picks import seg2picks @@ -32,4 +35,8 @@ "seg2picks", "mesh2seg", "seg2mesh", + # Filament commands + "seg2fil", + "fil2picks", + "fil2seg", ] diff --git a/src/copick_utils/cli/fil2picks.py b/src/copick_utils/cli/fil2picks.py new file mode 100644 index 0000000..458f9d3 --- /dev/null +++ b/src/copick_utils/cli/fil2picks.py @@ -0,0 +1,135 @@ +"""CLI command for sampling picks along filaments.""" + +import click +import copick +from click_option_group import optgroup +from copick.cli.util import add_config_option, add_debug_option, add_run_names_option +from copick.util.log import get_logger +from copick.util.uri import parse_copick_uri + +from copick_utils.cli.util import ( + add_filament_sampling_options, + add_input_option, + add_output_option, + add_workers_option, +) +from copick_utils.util.config_models import create_simple_config + + +@click.command( + context_settings={"show_default": True}, + short_help="Sample picks along filaments.", + no_args_is_help=True, +) +@add_config_option +@add_run_names_option +@optgroup.group("\nInput Options", help="Options related to the input filaments.") +@add_input_option("filaments") +@optgroup.group("\nTool Options", help="Options related to this tool.") +@add_filament_sampling_options +@add_workers_option +@optgroup.group("\nOutput Options", help="Options related to the output picks.") +@add_output_option("picks", default_tool="fil2picks") +@add_debug_option +def fil2picks( + config, + run_names, + input_uri, + spacing, + anchor, + roll, + seed, + length_unit, + workers, + output_uri, + debug, +): + """ + Sample picks along filaments. + + Places picks along each filament at a fixed `--spacing`, measured along the filament's + curve. Each pick's orientation has its +Z axis along the filament, in the filament's point + order, and its rotation about that axis follows `--roll`. Every pick carries its filament's + ID as its instance ID and the filament's score, and the picks are written filament by + filament, in order along each one. Exported to RELION with + `copick export picks --filament-columns on`, they carry RELION's filament conventions. + + A filament's stored curve is used when it still matches the filament: the B-spline that + `copick convert seg2fil` fitted, or control points from an editor. A filament without one is + sampled along a smooth curve derived from its points. Sampling is cheap and does not change + the filaments, so different spacings can be tried without tracing again. + + URI Format: + + \b + Filaments: object_name:user_id/session_id + Picks: object_name:user_id/session_id + + Examples: + + \b + # Picks every 82 Å (one tubulin dimer), in the filaments' own session + copick convert fil2picks -i "microtubule:trace/1" -o "microtubule:trace/1" --spacing 82 + + \b + # Every 4 nm with a random rotation about the axis + copick convert fil2picks -i "microtubule:trace/1" -o "microtubule:picks/4nm" --spacing 40 --roll random --seed 1 + + See Also: + + \b + copick convert seg2fil: trace filaments in a segmentation + copick convert fil2seg: paint tubes around filaments into an instance segmentation + copick export picks: export picks to RELION with --filament-columns + """ + from copick_utils.converters.picks_from_filaments import picks_from_filaments_lazy_batch + + logger = get_logger(__name__, debug=debug) + + root = copick.from_file(config) + run_names_list = list(run_names) if run_names else None + + try: + task_config = create_simple_config( + input_uri=input_uri, + input_type="filaments", + output_uri=output_uri, + output_type="picks", + command_name="fil2picks", + ) + except ValueError as e: + raise click.BadParameter(str(e)) from e + + input_params = parse_copick_uri(input_uri, "filaments") + logger.info(f"Sampling filaments '{input_params['object_name']}' every {spacing} ({length_unit})") + + results = picks_from_filaments_lazy_batch( + root=root, + config=task_config, + run_names=run_names_list, + workers=workers, + spacing=spacing, + anchor=anchor, + roll=roll, + seed=seed, + length_unit=length_unit, + ) + + successful = sum(1 for result in results.values() if result and result.get("processed", 0) > 0) + filaments = sum(result.get("filaments_sampled", 0) for result in results.values() if result) + points = sum(result.get("points_created", 0) for result in results.values() if result) + + all_errors = [] + for result in results.values(): + if result and result.get("errors"): + all_errors.extend(result["errors"]) + + logger.info(f"Completed: {successful}/{len(results)} runs processed successfully") + logger.info(f"Total: {points} picks from {filaments} filaments") + + if all_errors: + logger.warning(f"Encountered {len(all_errors)} errors during processing") + for error in all_errors[:5]: + logger.warning(f" - {error}") + if len(all_errors) > 5: + logger.warning(f" ... and {len(all_errors) - 5} more errors") diff --git a/src/copick_utils/cli/fil2seg.py b/src/copick_utils/cli/fil2seg.py new file mode 100644 index 0000000..0b1a617 --- /dev/null +++ b/src/copick_utils/cli/fil2seg.py @@ -0,0 +1,139 @@ +"""CLI command for painting tubes around filaments into an instance segmentation.""" + +import click +import copick +from click_option_group import optgroup +from copick.cli.util import add_config_option, add_debug_option, add_run_names_option +from copick.util.log import get_logger +from copick.util.uri import parse_copick_uri + +from copick_utils.cli.util import add_input_option, add_output_option, add_workers_option +from copick_utils.util.config_models import create_simple_config + + +@click.command( + context_settings={"show_default": True}, + short_help="Paint tubes around filaments into an instance segmentation.", + no_args_is_help=True, +) +@add_config_option +@add_run_names_option +@optgroup.group("\nInput Options", help="Options related to the input filaments.") +@add_input_option("filaments") +@optgroup.group("\nTool Options", help="Options related to this tool.") +@optgroup.option( + "--radius", + type=float, + default=None, + help="Tube radius in angstroms for filaments without their own radius. Unset: the object's radius.", +) +@optgroup.option( + "--tomo-type", + "-tt", + default="wbp", + help="Type of tomogram whose shape the segmentation takes.", +) +@add_workers_option +@optgroup.group("\nOutput Options", help="Options related to the output segmentation.") +@add_output_option("segmentation", default_tool="fil2seg") +@add_debug_option +def fil2seg( + config, + run_names, + input_uri, + radius, + tomo_type, + workers, + output_uri, + debug, +): + """ + Paint tubes around filaments into an instance segmentation. + + Each filament becomes a tube around its centreline, labelled with the filament's ID, so the + result is an instance segmentation with the same IDs as the filaments and their picks. + A filament's tube radius is its own radius (`copick convert seg2fil` records the radius of + the segmentation it traced), otherwise `--radius`, otherwise the object's radius. Where + tubes overlap, each voxel goes to the nearest centreline. The segmentation has the shape of + the run's `--tomo-type` tomogram at the voxel spacing given in the output URI. + + Useful for filaments drawn or edited by hand, or imported, that have no segmentation yet: + for example as per-filament masks, or as input to `copick convert seg2fil`. + + URI Format: + + \b + Filaments: object_name:user_id/session_id + Segmentations: name:user_id/session_id@voxel_spacing?instance=true + + Examples: + + \b + # Tubes of each filament's own radius, at 10 Å + copick convert fil2seg -i "microtubule:manual/1" -o "microtubule:tubes/1@10.0?instance=true" + + \b + # 125 Å tubes against a denoised tomogram + copick convert fil2seg -i "microtubule:manual/1" -o "microtubule:tubes/1@10.0?instance=true" \\ + --radius 125 --tomo-type denoised + + See Also: + + \b + copick convert seg2fil: trace filaments in a segmentation + copick convert fil2picks: sample picks along filaments + copick convert picks2seg: paint spheres around picks + """ + from copick_utils.converters.segmentation_from_filaments import segmentation_from_filaments_lazy_batch + + logger = get_logger(__name__, debug=debug) + + root = copick.from_file(config) + run_names_list = list(run_names) if run_names else None + + output_params = parse_copick_uri(output_uri, "segmentation") + if output_params.get("multilabel") or output_params.get("panoptic"): + raise click.BadParameter("fil2seg writes an instance segmentation (?instance=true).") + if not output_params.get("instance"): + output_uri = output_uri + ("&" if "?" in output_uri else "?") + "instance=true" + + try: + task_config = create_simple_config( + input_uri=input_uri, + input_type="filaments", + output_uri=output_uri, + output_type="segmentation", + command_name="fil2seg", + ) + except ValueError as e: + raise click.BadParameter(str(e)) from e + + input_params = parse_copick_uri(input_uri, "filaments") + logger.info(f"Painting filaments '{input_params['object_name']}' into an instance segmentation") + + results = segmentation_from_filaments_lazy_batch( + root=root, + config=task_config, + run_names=run_names_list, + workers=workers, + radius=radius, + tomo_type=tomo_type, + ) + + successful = sum(1 for result in results.values() if result and result.get("processed", 0) > 0) + filaments = sum(result.get("filaments_painted", 0) for result in results.values() if result) + + all_errors = [] + for result in results.values(): + if result and result.get("errors"): + all_errors.extend(result["errors"]) + + logger.info(f"Completed: {successful}/{len(results)} runs processed successfully") + logger.info(f"Total filaments painted: {filaments}") + + if all_errors: + logger.warning(f"Encountered {len(all_errors)} errors during processing") + for error in all_errors[:5]: + logger.warning(f" - {error}") + if len(all_errors) > 5: + logger.warning(f" ... and {len(all_errors) - 5} more errors") diff --git a/src/copick_utils/cli/fit_spline.py b/src/copick_utils/cli/fit_spline.py index 03961e6..87962b7 100644 --- a/src/copick_utils/cli/fit_spline.py +++ b/src/copick_utils/cli/fit_spline.py @@ -8,9 +8,9 @@ resolve_deprecated_option, ) from copick.util.log import get_logger -from copick.util.uri import parse_copick_uri +from copick.util.uri import expand_output_uri, parse_copick_uri -from copick_utils.cli.util import add_input_option, add_output_option +from copick_utils.cli.util import add_input_option, add_output_option, add_workers_option from copick_utils.util.config_models import create_simple_config @@ -38,12 +38,12 @@ "--spacing-distance", type=float, required=True, - help="Distance between consecutive sampled points along the spline.", + help="Distance between consecutive sampled points along each spline, in voxels.", ) @optgroup.option( "--smoothing-factor", type=float, - help="Smoothing parameter for spline fitting (auto if not provided).", + help="Smoothing parameter for spline fitting (scipy's s, in voxels squared; auto if not provided).", ) @optgroup.option( "--degree", @@ -51,11 +51,13 @@ default=3, help="Degree of the spline (1-5).", ) +# TODO:remove once deprecation takes effect -- --connectivity-radius has no effect (26-connected skeleton graph) @optgroup.option( "--connectivity-radius", type=float, - default=2.0, - help="Maximum distance to consider skeleton points as connected.", + default=None, + hidden=True, + help="Deprecated and ignored: skeleton voxels are joined to their 26 neighbours.", ) @optgroup.option( "--compute-transforms/--no-compute-transforms", @@ -67,22 +69,32 @@ "--curvature-threshold", type=float, default=0.2, - help="Maximum allowed curvature before outlier removal.", + help="Largest sine of the turning angle between consecutive sampled points; a spline that turns more " + "sharply is smoothed further.", ) @optgroup.option( "--max-iterations", type=int, default=5, - help="Maximum number of outlier removal iterations.", + help="Maximum number of smoothing increases.", ) @optgroup.option( - "--workers", + "--label", type=int, - default=8, - help="Number of worker processes.", + default=None, + help="Label to fit in a multilabel segmentation (default: every non-zero voxel).", ) +@add_workers_option @optgroup.group("\nOutput Options", help="Options related to output picks.") @add_output_option("picks", default_tool="spline") +@add_output_option( + "filaments", + flag="--filaments", + short_flag="-of", + param_name="filaments_uri", + required=False, + description="Also store the fitted splines as copick Filaments (exact B-spline curves) under this URI.", +) @add_debug_option def fit_spline( config, @@ -96,19 +108,30 @@ def fit_spline( compute_transforms, curvature_threshold, max_iterations, + label, workers, output_uri, + filaments_uri, debug, ): """Fit 3D splines to skeletons and generate oriented picks. - Fits regularized 3D parametric splines to skeletonized segmentation volumes and samples - points along each spline at a regular interval, producing picks. Orientations are computed - from the local spline direction when `--compute-transforms` is enabled. + Fits a smoothing 3D spline to every filament of a skeleton (or segmentation) volume and + samples points along each spline exactly `--spacing-distance` voxels apart, producing picks. + The skeleton is split into filaments between ends and junctions, continuing straight + through junctions, so crossing or separate filaments each get their own spline. All picks go + into one pick set: grouped by filament, in order along it, with the filament's number + (1, 2, ... by length) as the pick's instance ID. With `--compute-transforms`, each pick's + +Z axis follows the spline's direction, and the rotation about it changes as little as + possible along the filament. + + Where a spline turns more sharply than `--curvature-threshold` between samples, its + smoothing is increased (up to `--max-iterations` times). `--filaments` also stores the + fitted splines as copick Filaments. - Curvature-based outlier removal is applied iteratively to discard skeleton points that - produce unrealistically sharp bends, while the connectivity radius controls how skeleton - voxels are joined into a connected curve before fitting. + For new work, `copick convert seg2fil` (tracing, with an instance segmentation and + thickness-based filters) and `copick convert fil2picks` (sampling in angstroms) separate + tracing from sampling. URI Format: @@ -120,23 +143,27 @@ def fit_spline( \b # Fit splines to skeletonized components (voxel spacing from the @10.0 in -i) - copick process fit_spline -i "skeleton:skel/inst-.*@10.0" \\ + copick process fit-spline -i "skeleton:skel/inst-.*@10.0" \\ -o "skeleton:spline/spline-{input_session_id}" --spacing-distance 4.4 \b # Process a single skeleton component - copick process fit_spline -i "skeleton:skel/skel-0@10.0" \\ + copick process fit-spline -i "skeleton:skel/skel-0@10.0" \\ -o "skeleton:spline/spline-0" --spacing-distance 2.0 See Also: \b + copick convert seg2fil: trace filaments in a segmentation (Filaments and an instance segmentation) + copick convert fil2picks: sample picks along filaments, in angstroms copick process skeletonize: produce the skeleton segmentations fed to this command - copick process separate_components: split a segmentation into per-instance skeletons first """ from copick_utils.process.spline_fitting import fit_spline_lazy_batch logger = get_logger(__name__, debug=debug) + logger.info("fit-spline is kept for compatibility; see copick convert seg2fil and fil2picks for new work") + if connectivity_radius is not None: + logger.warning("--connectivity-radius is deprecated and has no effect") root = copick.from_file(config) run_names_list = list(run_names) if run_names else None @@ -173,12 +200,30 @@ def fit_spline( "Input URI must include a specific voxel spacing (e.g., @10.0), or pass --voxel-spacing/-vs.", ) + filament_params = {} + if filaments_uri: + parsed = parse_copick_uri( + expand_output_uri( + output_uri=filaments_uri, + input_uri=input_uri, + input_type="segmentation", + output_type="filaments", + command_name="spline", + ), + "filaments", + ) + filament_params = { + "filaments_object_name": parsed["object_name"], + "filaments_user_id": parsed["user_id"], + "filaments_session_id": parsed["session_id"], + } + logger.info(f"Fitting splines to segmentations '{input_params['name']}'") logger.info( f"Source segmentation pattern: {input_params['name']} ({input_params['user_id']}/{input_params['session_id']})", ) logger.info(f"Spacing distance: {spacing_distance}, degree: {degree}") - logger.info(f"Smoothing factor: {smoothing_factor}, connectivity radius: {connectivity_radius}") + logger.info(f"Smoothing factor: {smoothing_factor}") logger.info(f"Compute transforms: {compute_transforms}") logger.info(f"Curvature threshold: {curvature_threshold}, max iterations: {max_iterations}") logger.info(f"Voxel spacing: {voxel_spacing}") @@ -193,11 +238,12 @@ def fit_spline( spacing_distance=spacing_distance, smoothing_factor=smoothing_factor, degree=degree, - connectivity_radius=connectivity_radius, compute_transforms=compute_transforms, curvature_threshold=curvature_threshold, max_iterations=max_iterations, voxel_spacing=voxel_spacing, + label=label, + **filament_params, ) successful = sum(1 for result in results.values() if result and result.get("processed", 0) > 0) @@ -212,7 +258,8 @@ def fit_spline( logger.info(f"Completed: {successful}/{len(results)} runs processed successfully") logger.info(f"Total conversion tasks completed: {total_processed}") - logger.info(f"Total picks created: {total_picks}") + total_filaments = sum(result.get("filaments_fitted", 0) for result in results.values() if result) + logger.info(f"Total picks created: {total_picks} on {total_filaments} filaments") if all_errors: logger.warning(f"Encountered {len(all_errors)} errors during processing") diff --git a/src/copick_utils/cli/seg2fil.py b/src/copick_utils/cli/seg2fil.py new file mode 100644 index 0000000..d7d5c95 --- /dev/null +++ b/src/copick_utils/cli/seg2fil.py @@ -0,0 +1,226 @@ +"""CLI command for tracing filaments in segmentations.""" + +import click +import copick +from click_option_group import optgroup +from copick.cli.util import add_config_option, add_debug_option, add_run_names_option +from copick.util.log import get_logger +from copick.util.uri import expand_output_uri, parse_copick_uri + +from copick_utils.cli.util import ( + add_filament_tracing_options, + add_input_option, + add_output_option, + add_workers_option, +) +from copick_utils.util.config_models import create_simple_config + + +def length_in_angstrom(value, unit, input_uri): + """A length option in angstroms, converted from voxels with the input URI's voxel spacing if needed.""" + if value is None or unit == "angstrom": + return value + vs_raw = parse_copick_uri(input_uri, "segmentation").get("voxel_spacing") + if vs_raw is None or vs_raw == "*": + raise click.BadParameter("Lengths in voxels require a voxel spacing in the input URI (e.g., @10.0).") + return value * float(vs_raw) + + +def volume_in_angstrom3(value, unit, input_uri): + """A volume option in cubic angstroms, converted from cubic voxels with the input URI's voxel spacing if needed.""" + if value is None or unit == "angstrom": + return value + vs_raw = parse_copick_uri(input_uri, "segmentation").get("voxel_spacing") + if vs_raw is None or vs_raw == "*": + raise click.BadParameter("Volumes in voxels require a voxel spacing in the input URI (e.g., @10.0).") + return value * float(vs_raw) ** 3 + + +@click.command( + context_settings={"show_default": True}, + short_help="Trace filaments in segmentations.", + no_args_is_help=True, +) +@add_config_option +@add_run_names_option +@optgroup.group("\nInput Options", help="Options related to the input segmentation.") +@add_input_option("segmentation") +@optgroup.group("\nTool Options", help="Options related to this tool.") +@add_filament_tracing_options +@add_workers_option +@optgroup.group("\nOutput Options", help="Options related to the output filaments.") +@add_output_option("filaments", default_tool="seg2fil") +@add_output_option( + "segmentation", + flag="--instances", + short_flag="-oi", + param_name="instances_uri", + required=False, + description="Also write the instance segmentation of the traced filaments (each voxel holds the ID of its " + "filament, the same IDs as the filaments). Smart defaults as for -o; ?instance=true is implied.", +) +@add_debug_option +def seg2fil( + config, + run_names, + input_uri, + min_length, + min_aspect, + min_radius, + fill_lumen, + min_volume, + prune_length, + junction_merge, + max_bend, + smoothing, + extend_ends, + label, + length_unit, + volume_unit, + workers, + output_uri, + instances_uri, + debug, +): + """ + Trace filaments in segmentations. + + Traces the centreline of every filament (e.g. microtubules or actin) in a segmentation and + stores it as a copick Filaments entry: the filament's fitted B-spline (an exact `bspline` + curve that editors can reopen) and the centreline points copick regenerates from it. + Filaments are numbered 1, 2, ... by length, longest first. + + Each connected component is skeletonized and split into branches between ends and junctions. + Holes up to `--fill-lumen` are filled first, so a tube labelled by its wall alone traces as one + filament. Side branches shorter than `--prune-length` are pruned, and at each junction the branches + that continue most nearly straight (within `--max-bend`) are joined, so two crossing + filaments stay two filaments. Free ends are extended to the edge of the segmentation, and + filaments shorter than `--min-length`, shorter than `--min-aspect` label diameters (blobs), + or thinner than `--min-radius` (slivers of segmentation noise) are rejected. Options left + unset are derived from the segmentation's own thickness or the object's radius. + + The input can be a binary, multilabel (`--label`), panoptic or instance segmentation. An + instance segmentation, for example one edited by hand or written by `--instances`, is traced + one instance at a time and keeps its IDs. Sample picks from the filaments with + `copick convert fil2picks`. + + URI Format: + + \b + Segmentations: name:user_id/session_id@voxel_spacing + Instance or panoptic input: append ?instance=true or ?panoptic=true + Filaments: object_name:user_id/session_id + + Examples: + + \b + # Trace microtubules and keep the instance segmentation of the traced filaments + copick convert seg2fil -i "microtubule:easymode/job006@10.0" \\ + -o "microtubule:trace/1" --instances "microtubule:trace/1@10.0?instance=true" + + \b + # Reject filaments shorter than 100 nm + copick convert seg2fil -i "microtubule:easymode/job006@10.0" -o "microtubule:trace/1" --min-length 1000 + + \b + # Re-trace an edited instance segmentation, keeping its IDs + copick convert seg2fil -i "microtubule:curated/1@10.0?instance=true" -o "microtubule:trace/curated" + + See Also: + + \b + copick convert fil2picks: sample picks along the traced filaments + copick convert fil2seg: paint tubes around filaments into an instance segmentation + copick process fit-spline: segmentation to picks in one step (older, voxel units) + copick process seg-stats: component and skeleton statistics to choose thresholds + """ + from copick_utils.converters.filaments_from_segmentation import filaments_from_segmentation_lazy_batch + + logger = get_logger(__name__, debug=debug) + + root = copick.from_file(config) + run_names_list = list(run_names) if run_names else None + + min_length = length_in_angstrom(min_length, length_unit, input_uri) + min_radius = length_in_angstrom(min_radius, length_unit, input_uri) + fill_lumen = length_in_angstrom(fill_lumen, length_unit, input_uri) + prune_length = length_in_angstrom(prune_length, length_unit, input_uri) + junction_merge = length_in_angstrom(junction_merge, length_unit, input_uri) + smoothing = length_in_angstrom(smoothing, length_unit, input_uri) + min_volume = volume_in_angstrom3(min_volume, volume_unit, input_uri) + + try: + task_config = create_simple_config( + input_uri=input_uri, + input_type="segmentation", + output_uri=output_uri, + output_type="filaments", + command_name="seg2fil", + ) + except ValueError as e: + raise click.BadParameter(str(e)) from e + + instance_params = {} + if instances_uri: + expanded = expand_output_uri( + output_uri=instances_uri, + input_uri=input_uri, + input_type="segmentation", + output_type="segmentation", + command_name="seg2fil", + ) + parsed = parse_copick_uri(expanded, "segmentation") + if parsed.get("multilabel") or parsed.get("panoptic"): + raise click.BadParameter("--instances writes an instance segmentation (?instance=true).") + instance_params = { + "instances_object_name": parsed["name"], + "instances_user_id": parsed["user_id"], + "instances_session_id": parsed["session_id"], + } + + input_params = parse_copick_uri(input_uri, "segmentation") + logger.info(f"Tracing filaments in segmentation '{input_params['name']}'") + logger.info(f"Source segmentation pattern: {input_params['user_id']}/{input_params['session_id']}") + + results = filaments_from_segmentation_lazy_batch( + root=root, + config=task_config, + run_names=run_names_list, + workers=workers, + min_volume=min_volume, + min_length=min_length, + min_aspect=min_aspect, + min_radius=min_radius, + fill_lumen=fill_lumen, + prune_length=prune_length, + junction_merge=junction_merge, + max_bend=max_bend, + smoothing=smoothing, + extend_ends=extend_ends, + label=label, + **instance_params, + ) + + successful = sum(1 for result in results.values() if result and result.get("processed", 0) > 0) + filaments = sum(result.get("filaments_written", 0) for result in results.values() if result) + length = sum(result.get("total_length", 0.0) for result in results.values() if result) + short = sum(result.get("rejected_length", 0) for result in results.values() if result) + blobs = sum(result.get("rejected_aspect", 0) for result in results.values() if result) + thin = sum(result.get("rejected_radius", 0) for result in results.values() if result) + + all_errors = [] + for result in results.values(): + if result and result.get("errors"): + all_errors.extend(result["errors"]) + + logger.info(f"Completed: {successful}/{len(results)} runs processed successfully") + logger.info( + f"Total filaments: {filaments} ({length / 10:.0f} nm); rejected {short} short, {blobs} blob-like, {thin} thin", + ) + + if all_errors: + logger.warning(f"Encountered {len(all_errors)} errors during processing") + for error in all_errors[:5]: + logger.warning(f" - {error}") + if len(all_errors) > 5: + logger.warning(f" ... and {len(all_errors) - 5} more errors") diff --git a/src/copick_utils/cli/skeletonize.py b/src/copick_utils/cli/skeletonize.py index ec068a8..f54405e 100644 --- a/src/copick_utils/cli/skeletonize.py +++ b/src/copick_utils/cli/skeletonize.py @@ -41,13 +41,26 @@ "--remove-short-branches/--keep-short-branches", is_flag=True, default=True, - help="Remove short branches from skeleton.", + help="Remove small skeleton pieces (whole connected pieces with fewer than --min-branch-length voxels). " + "To shorten side branches of a larger skeleton, use --prune-length.", ) @optgroup.option( "--min-branch-length", type=int, default=5, - help="Minimum length of branches to keep.", + help="Minimum number of voxels of a skeleton piece to keep (see --remove-short-branches).", +) +@optgroup.option( + "--prune-length", + type=float, + default=None, + help="Prune side branches (spurs) of the skeleton shorter than this (unit: --length-unit). Unset: no pruning.", +) +@optgroup.option( + "--length-unit", + type=click.Choice(["angstrom", "voxel"]), + default="angstrom", + help="Unit of --prune-length.", ) @add_workers_option @optgroup.group("\nOutput Options", help="Options related to output segmentations.") @@ -62,6 +75,8 @@ def skeletonize( min_object_size, remove_short_branches, min_branch_length, + prune_length, + length_unit, workers, output_uri, debug, @@ -76,7 +91,9 @@ def skeletonize( The input session ID is treated as a regex, so a single invocation can skeletonize many segmentations at once. This pairs naturally with the output of connected-component separation (e.g. pattern `inst-.*` to match `inst-0`, `inst-1`, etc.). Optional cleanup - removes small objects before thinning and prunes short spur branches afterwards. + removes small objects before thinning and small skeleton pieces afterwards, and + `--prune-length` prunes short side branches (spurs) off the skeleton. The output takes the + name given in `-o`. URI Format: @@ -94,6 +111,11 @@ def skeletonize( copick process skeletonize -i "membrane:user1/inst-.*@10.0" \\ -o "membrane:skel/skel-{input_session_id}@10.0" + \b + # Prune side branches shorter than 20 nm + copick process skeletonize -i "microtubule:easymode/job006@10.0" -o "microtubule:skel/job006@10.0" \\ + --prune-length 200 + \b # Use the distance-transform backend and keep short branches copick process skeletonize --method distance_transform --keep-short-branches \\ @@ -107,6 +129,12 @@ def skeletonize( """ from copick_utils.process.skeletonize import skeletonize_lazy_batch + if prune_length is not None and length_unit == "voxel": + vs_raw = parse_copick_uri(input_uri, "segmentation").get("voxel_spacing") + if vs_raw is None or vs_raw == "*": + raise click.BadParameter("--length-unit voxel requires voxel spacing in the input URI (e.g., @10.0)") + prune_length = prune_length * float(vs_raw) + logger = get_logger(__name__, debug=debug) root = copick.from_file(config) @@ -143,6 +171,7 @@ def skeletonize( min_object_size=min_object_size, remove_short_branches=remove_short_branches, min_branch_length=min_branch_length, + prune_length=prune_length, ) successful = sum(1 for result in results.values() if result and result.get("processed", 0) > 0) diff --git a/src/copick_utils/cli/util.py b/src/copick_utils/cli/util.py index d611d1e..05e78eb 100644 --- a/src/copick_utils/cli/util.py +++ b/src/copick_utils/cli/util.py @@ -739,3 +739,175 @@ def add_invert_option_decorator(_func: click.Command) -> click.Command: return add_invert_option_decorator else: return add_invert_option_decorator(func) + + +def add_filament_tracing_options(func: click.Command = None) -> Callable: + """ + Add the filament tracing options (seg2fil, fit-spline). + + Lengths are in angstroms unless --length-unit voxel is given, and --min-volume in cubic angstroms unless + --volume-unit voxel is given. Options left unset take values derived from the segmentation itself, so they carry + over between objects and voxel sizes. + + Args: + func (click.Command, optional): The Click command to which the options will be added. + + Returns: + Callable: The Click command with the tracing options added. + """ + + def add_filament_tracing_options_decorator(_func: click.Command) -> click.Command: + opts = [ + optgroup.option( + "--min-length", + type=float, + default=None, + help="Reject filaments shorter than this (unit: --length-unit). Unset: no length filter.", + ), + optgroup.option( + "--min-aspect", + type=float, + default=3.0, + help="Reject filaments shorter than this many label diameters (compact blobs and specks). The " + "default rejects only microtubule-thick stubs under ~44 nm on dataset 10521.", + ), + optgroup.option( + "--min-radius", + type=float, + default=None, + help="Reject filaments whose median label radius is below this (unit: --length-unit): thin slivers " + "of segmentation noise. Unset: a third of the object's (tube) radius.", + ), + optgroup.option( + "--fill-lumen", + type=float, + default=None, + help="Before tracing, fill holes up to this radius (unit: --length-unit) in every slice of the label, " + "so a tube whose wall alone is labelled traces as one filament. Unset: the object's (tube) radius; " + "0: no filling.", + ), + optgroup.option( + "--min-volume", + type=float, + default=None, + help="Drop connected components smaller than this before tracing (unit: --volume-unit). " + "Unset: no volume filter.", + ), + optgroup.option( + "--prune-length", + type=float, + default=None, + help="Prune skeleton side branches shorter than this (unit: --length-unit). " + "Unset: one label diameter.", + ), + optgroup.option( + "--junction-merge", + type=float, + default=None, + help="Merge junctions joined by a bridge up to this long, as two crossing filaments often " + "skeletonize (unit: --length-unit). Unset: one label diameter.", + ), + optgroup.option( + "--max-bend", + type=float, + default=45.0, + help="Largest deviation from a straight line, in degrees, for a filament to continue through a " + "junction.", + ), + optgroup.option( + "--smoothing", + type=float, + default=None, + help="RMS deviation of the fitted spline from the skeleton (unit: --length-unit). " + "Unset: half a voxel.", + ), + optgroup.option( + "--extend-ends/--no-extend-ends", + default=True, + help="Extend free filament ends along their direction to the edge of the segmentation (thinning " + "shortens each end by about one radius).", + ), + optgroup.option( + "--label", + type=int, + default=None, + help="Label to trace in a multilabel or panoptic segmentation (default: the output object's label).", + ), + optgroup.option( + "--length-unit", + type=click.Choice(["angstrom", "voxel"]), + default="angstrom", + help="Unit of --min-length, --min-radius, --fill-lumen, --prune-length, --junction-merge and --smoothing.", + ), + optgroup.option( + "--volume-unit", + type=click.Choice(["angstrom", "voxel"]), + default="angstrom", + help="Unit of --min-volume: 'angstrom' for ų, 'voxel' for cubic voxels.", + ), + ] + for opt in reversed(opts): + _func = opt(_func) + return _func + + if func is None: + return add_filament_tracing_options_decorator + else: + return add_filament_tracing_options_decorator(func) + + +def add_filament_sampling_options(func: click.Command = None) -> Callable: + """ + Add the filament sampling options (fil2picks). + + Args: + func (click.Command, optional): The Click command to which the options will be added. + + Returns: + Callable: The Click command with the sampling options added. + """ + + def add_filament_sampling_options_decorator(_func: click.Command) -> click.Command: + opts = [ + optgroup.option( + "--spacing", + type=float, + required=True, + help="Distance between picks along a filament (unit: --length-unit). Required: no spacing suits " + "every filament.", + ), + optgroup.option( + "--anchor", + type=click.Choice(["center", "start"]), + default="center", + help="'center' splits the length left after the last full step evenly between both ends; " + "'start' places the first pick at the filament's start.", + ), + optgroup.option( + "--roll", + type=click.Choice(["parallel", "random"]), + default="parallel", + help="Rotation of the picks about the filament axis: 'parallel' turns as little as possible along " + "the filament (rotation-minimizing frames); 'random' is uniformly random per pick.", + ), + optgroup.option( + "--seed", + type=int, + default=None, + help="Random seed for --roll random.", + ), + optgroup.option( + "--length-unit", + type=click.Choice(["angstrom", "voxel"]), + default="angstrom", + help="Unit of --spacing ('voxel' needs the filaments' voxel spacing).", + ), + ] + for opt in reversed(opts): + _func = opt(_func) + return _func + + if func is None: + return add_filament_sampling_options_decorator + else: + return add_filament_sampling_options_decorator(func) diff --git a/src/copick_utils/converters/filaments_from_segmentation.py b/src/copick_utils/converters/filaments_from_segmentation.py new file mode 100644 index 0000000..a46dda4 --- /dev/null +++ b/src/copick_utils/converters/filaments_from_segmentation.py @@ -0,0 +1,197 @@ +"""Trace filaments in segmentations: copick Filaments (fitted B-spline curves) and an instance segmentation.""" + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import numpy as np +from copick.util.log import get_logger + +from copick_utils.converters.lazy_converter import create_lazy_batch_converter +from copick_utils.process.filament_tracing import TraceParameters, trace_centrelines + +if TYPE_CHECKING: + from copick.models import CopickFilaments, CopickRun, CopickSegmentation + +logger = get_logger(__name__) + +#: Default minimum label radius, as a fraction of the filament object's tube radius (see ``min_radius``). +MIN_RADIUS_FRACTION = 1.0 / 3.0 + + +def _volume_to_trace(segmentation: "CopickSegmentation", object_name: str, label: Optional[int]): + """The volume to trace and whether it holds instances, from the segmentation's type.""" + from copick_utils.util.segmentations import segmentation_type + + seg_type = segmentation_type(segmentation) + if seg_type == "instance": + return segmentation.numpy(), True + if seg_type == "binary": + return segmentation.numpy() > 0, False + + if label is None: + obj = segmentation.run.root.get_object(object_name) + if obj is None or obj.label is None: + raise ValueError( + f"Choose the label to trace in the {seg_type} segmentation {segmentation.name} with --label " + f"({object_name} is not a pickable object with a label).", + ) + label = int(obj.label) + if seg_type == "multilabel": + return segmentation.numpy() == label, False + + # Panoptic: the object's instances if it has any, otherwise its label region + channels = segmentation.numpy() + labels, instances = channels[0], channels[1] + in_label = labels == label + if np.any(instances[in_label] > 0): + return np.where(in_label, instances, 0), True + return in_label, False + + +def filaments_from_segmentation( + segmentation: "CopickSegmentation", + run: "CopickRun", + object_name: str, + session_id: str, + user_id: str, + min_volume: Optional[float] = None, + min_length: Optional[float] = None, + min_aspect: float = 3.0, + min_radius: Optional[float] = None, + prune_length: Optional[float] = None, + junction_merge: Optional[float] = None, + max_bend: float = 45.0, + smoothing: Optional[float] = None, + extend_ends: bool = True, + fill_lumen: Optional[float] = None, + label: Optional[int] = None, + instances_object_name: Optional[str] = None, + instances_user_id: Optional[str] = None, + instances_session_id: Optional[str] = None, + **kwargs, +) -> Optional[Tuple["CopickFilaments", Dict[str, float]]]: + """ + Trace the filaments of a segmentation and store them as copick Filaments. + + Each filament is stored with its fitted spline as a ``bspline`` curve (copick regenerates the filament's points + from it) and the fit's settings in ``metadata["fit"]``. IDs are 1..K by length (longest first) for a binary, + multilabel or panoptic-region input; an instance segmentation (or a panoptic segmentation's instances) keeps + its IDs. With ``instances_*`` set, the instance segmentation of the traced filaments is stored as well: each + label voxel holds the ID of the nearest filament of its connected component, with the same IDs as the + filaments. + + Args: + segmentation: Input segmentation (binary, multilabel, instance or panoptic). + run: CopickRun object. + object_name: Object name of the output filaments (a filament object). + session_id: Session ID of the output filaments. + user_id: User ID of the output filaments. + min_volume: Drop components smaller than this (ų). None: no volume filter. + min_length: Reject filaments shorter than this (Å). None: no length filter. + min_aspect: Reject filaments shorter than this many label diameters. + min_radius: Reject filaments whose label radius is below this (Å). None: a third of the object's radius + (the filament's tube radius), or no radius filter if the object has none. + prune_length: Prune skeleton spurs shorter than this (Å). None: one label diameter. + junction_merge: Merge junctions joined by a bridge up to this long (Å). None: one label diameter. + max_bend: Largest deviation from a straight line (degrees) to continue through a junction. + smoothing: RMS deviation of the fitted spline from the skeleton (Å). None: half a voxel. + extend_ends: Extend free ends to the edge of the label. + fill_lumen: Fill holes up to this radius (Å) in each axis-aligned slice before tracing, the empty lumen of + a tube whose wall alone is labelled. None: the object's radius; 0: no filling. + label: Label to trace in a multilabel or panoptic segmentation (default: the object's label). + instances_object_name: Object name of the instance segmentation to store, or None for none. + instances_user_id: User ID of the instance segmentation. + instances_session_id: Session ID of the instance segmentation ({input_session_id} is replaced). + **kwargs: Additional keyword arguments from the lazy converter. + + Returns: + Tuple of (CopickFilaments, stats dict) or None if the operation failed. + """ + from copick.models import CopickFilamentCurve + + import copick_utils + from copick_utils.util.segmentations import new_segmentation + + try: + voxel_size = float(segmentation.voxel_size) + volume, instances = _volume_to_trace(segmentation, object_name, label) + obj = run.root.get_object(object_name) + tube_radius = float(obj.radius) if obj is not None and obj.radius else None + if min_radius is None and tube_radius: + # Calibrated on dataset 10521 (easymode microtubule labels): real filaments have a label radius of + # 40-76 Å against a 120 Å tube radius; slivers and stubs of segmentation noise 10-35 Å. + min_radius = MIN_RADIUS_FRACTION * tube_radius + if fill_lumen is None: + fill_lumen = tube_radius # a hollow label's lumen is narrower than the tube + params = TraceParameters( + min_volume=min_volume, + min_length=min_length, + min_aspect=min_aspect, + min_radius=min_radius, + prune_length=prune_length, + junction_merge=junction_merge, + max_bend=max_bend, + smoothing=smoothing, + extend_ends=extend_ends, + fill_lumen=fill_lumen or None, + ) + centrelines, assignment, report = trace_centrelines(volume, voxel_size, params, instances=instances) + + source = f"{segmentation.name}:{segmentation.user_id}/{segmentation.session_id}@{voxel_size}" + curves, metadata = [], [] + for centreline in centrelines: + curves.append( + CopickFilamentCurve.from_tck( + centreline.tck, + step=voxel_size, + smoothing=centreline.smoothing, + scale=1.0, + ), + ) + fit = dict(centreline.metadata) + fit.update({"tool": "copick-utils seg2fil", "version": copick_utils.__version__, "source": source}) + metadata.append({"fit": fit}) + + filaments = run.new_filaments(object_name, session_id, user_id, exist_ok=True) + filaments.from_curves( + curves, + instance_ids=[c.instance_id for c in centrelines], + scores=[c.score for c in centrelines], + radii=[c.radius for c in centrelines], + voxel_spacing=voxel_size, + metadata=metadata, + ) + + stats = {key: value for key, value in report.as_dict().items() if isinstance(value, (int, float))} + stats["filaments_written"] = len(centrelines) + if instances_object_name and not instances: + instance_session = (instances_session_id or session_id).replace( + "{input_session_id}", + segmentation.session_id, + ) + output = new_segmentation( + run, + voxel_size, + instances_object_name, + instance_session, + instances_user_id or user_id, + "instance", + ) + output.from_numpy(assignment) + stats["instance_segmentations_written"] = 1 + + logger.info( + f"Traced {len(centrelines)} filaments ({report.total_length / 10:.0f} nm) in {run.name}; " + f"rejected {report.rejected_length} short, {report.rejected_aspect} blob-like and " + f"{report.rejected_radius} thin", + ) + return filaments, stats + + except Exception as e: + logger.error(f"Error tracing filaments in {run.name}: {e}") + return None + + +filaments_from_segmentation_lazy_batch = create_lazy_batch_converter( + converter_func=filaments_from_segmentation, + task_description="Tracing filaments", +) diff --git a/src/copick_utils/converters/picks_from_filaments.py b/src/copick_utils/converters/picks_from_filaments.py new file mode 100644 index 0000000..6a9b73c --- /dev/null +++ b/src/copick_utils/converters/picks_from_filaments.py @@ -0,0 +1,100 @@ +"""Sample picks along filaments at a fixed spacing, oriented along the filament and carrying its ID.""" + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import numpy as np +from copick.util.log import get_logger + +from copick_utils.converters.lazy_converter import create_lazy_batch_converter +from copick_utils.process.filament_tracing import sample_curve + +if TYPE_CHECKING: + from copick.models import CopickFilaments, CopickPicks, CopickRun + +logger = get_logger(__name__) + + +def picks_from_filaments( + filaments: "CopickFilaments", + run: "CopickRun", + object_name: str, + session_id: str, + user_id: str, + spacing: float, + anchor: str = "center", + roll: str = "parallel", + seed: Optional[int] = None, + length_unit: str = "angstrom", + **kwargs, +) -> Optional[Tuple["CopickPicks", Dict[str, int]]]: + """ + Sample each filament at a fixed spacing along its curve. + + A filament's stored curve is used when it still describes the filament (a traced ``bspline``, or ``catmull-rom`` + control points from an editor); otherwise copick derives a curve from its points. Each pick's rotation has its +Z + axis along the filament in point order, its translation is 0, and its ``instance_id`` and ``score`` are the + filament's. Picks are written grouped by filament and in order along it. + + Args: + filaments: Input CopickFilaments. + run: CopickRun object. + object_name: Object name of the output picks. + session_id: Session ID of the output picks. + user_id: User ID of the output picks. + spacing: Distance between picks along a filament, in Angstrom. Required: no spacing suits every filament. + anchor: ``center`` or ``start`` (see ``sample_curve``). + roll: ``parallel`` or ``random`` (see ``sample_curve``). + seed: Random seed for ``roll="random"`` (combined with each filament's ID). + length_unit: ``angstrom``, or ``voxel`` for a spacing in voxels of the filaments' voxel spacing. + **kwargs: Additional keyword arguments from the lazy converter. + + Returns: + Tuple of (CopickPicks, stats dict) or None if the operation failed. + """ + try: + if spacing is None or spacing <= 0: + raise ValueError("A positive spacing is required.") + if length_unit == "voxel": + if not filaments.voxel_spacing: + raise ValueError("A spacing in voxels needs the filaments' voxel spacing, which this file lacks.") + spacing = spacing * float(filaments.voxel_spacing) + positions, transforms, ids, scores = [], [], [], [] + for filament in filaments.filaments: + curve = filament.editable_curve().model_dump() + filament_seed = None if seed is None else int(seed) + int(filament.instance_id) + p, r, _ = sample_curve(curve, spacing, anchor=anchor, roll=roll, seed=filament_seed) + t = np.tile(np.eye(4), (len(p), 1, 1)) + t[:, :3, :3] = r + positions.append(p) + transforms.append(t) + ids.append(np.full(len(p), filament.instance_id, dtype=np.int64)) + scores.append(np.full(len(p), filament.score, dtype=float)) + + picks = run.new_picks(object_name, session_id, user_id, exist_ok=True) + if positions: + picks.from_numpy( + np.concatenate(positions), + np.concatenate(transforms), + instance_ids=np.concatenate(ids), + scores=np.concatenate(scores), + ) + else: + picks.points = [] + picks.store() + + stats = {"filaments_sampled": len(positions), "points_created": int(sum(len(p) for p in positions))} + logger.info( + f"Sampled {stats['points_created']} picks from {stats['filaments_sampled']} filaments " + f"every {spacing} Å in {run.name}", + ) + return picks, stats + + except Exception as e: + logger.error(f"Error sampling filaments in {run.name}: {e}") + return None + + +picks_from_filaments_lazy_batch = create_lazy_batch_converter( + converter_func=picks_from_filaments, + task_description="Sampling filaments", +) diff --git a/src/copick_utils/converters/segmentation_from_filaments.py b/src/copick_utils/converters/segmentation_from_filaments.py new file mode 100644 index 0000000..a863268 --- /dev/null +++ b/src/copick_utils/converters/segmentation_from_filaments.py @@ -0,0 +1,105 @@ +"""Paint tubes around filament centrelines: an instance segmentation whose IDs are the filaments' IDs.""" + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +import numpy as np +import zarr +from copick.util.log import get_logger + +from copick_utils.converters.lazy_converter import create_lazy_batch_converter +from copick_utils.process.filament_tracing import paint_tubes + +if TYPE_CHECKING: + from copick.models import CopickFilaments, CopickRun, CopickSegmentation + +logger = get_logger(__name__) + + +def _densify(points: np.ndarray, step: float) -> np.ndarray: + """Insert points along each segment so consecutive points are at most ``step`` apart.""" + points = np.asarray(points, dtype=float) + if len(points) < 2: + return points + out = [points[:1]] + for a, b in zip(points[:-1], points[1:]): + n = max(1, int(np.ceil(np.linalg.norm(b - a) / step))) + out.append(a + (b - a) * (np.arange(1, n + 1)[:, None] / n)) + return np.concatenate(out) + + +def segmentation_from_filaments( + filaments: "CopickFilaments", + run: "CopickRun", + object_name: str, + session_id: str, + user_id: str, + voxel_spacing: float, + tomo_type: str = "wbp", + radius: Optional[float] = None, + output_segmentation_type: Optional[str] = None, + **kwargs, +) -> Optional[Tuple["CopickSegmentation", Dict[str, int]]]: + """ + Paint a tube around each filament into an instance segmentation. + + Each filament's tube radius is its own ``radius`` if it has one, otherwise ``radius``, otherwise the object's + radius. Voxels inside a tube hold the filament's ID; where tubes overlap, the nearest centreline wins. The + volume's shape is that of the run's ``tomo_type`` tomogram at ``voxel_spacing``. + + Args: + filaments: Input CopickFilaments. + run: CopickRun object. + object_name: Object name of the output segmentation. + session_id: Session ID of the output segmentation. + user_id: User ID of the output segmentation. + voxel_spacing: Voxel spacing of the output segmentation, in Angstrom. + tomo_type: Tomogram type whose shape the output takes. + radius: Tube radius in Angstrom for filaments without their own radius. + output_segmentation_type: The type the output URI names; only an instance segmentation can be written. + **kwargs: Additional keyword arguments from the lazy converter. + + Returns: + Tuple of (CopickSegmentation, stats dict) or None if the operation failed. + """ + from copick_utils.util.segmentations import new_segmentation + + try: + if output_segmentation_type not in (None, "instance"): + raise ValueError(f"fil2seg writes instance segmentations, not {output_segmentation_type} ones.") + voxel_spacing = float(voxel_spacing) + vs = run.get_voxel_spacing(voxel_spacing) + tomograms = vs.get_tomograms(tomo_type) if vs is not None else [] + tomogram = tomograms[0] if tomograms else None + if tomogram is None: + raise ValueError(f"No {tomo_type} tomogram at {voxel_spacing} Å in {run.name}.") + shape = zarr.open(tomogram.zarr(), "r")["0"].shape + + obj = run.root.get_object(filaments.pickable_object_name) + fallback = radius if radius is not None else (obj.radius if obj is not None else None) + polylines, ids, radii = [], [], [] + for filament in filaments.filaments: + tube = filament.radius if filament.radius is not None else fallback + if tube is None: + raise ValueError( + f"Filament {filament.instance_id} has no radius; give one with --radius.", + ) + polylines.append(_densify(np.asarray(filament.points, dtype=float), voxel_spacing / 2.0)) + ids.append(int(filament.instance_id)) + radii.append(float(tube)) + + volume = paint_tubes(polylines, ids, radii, shape, voxel_spacing) + output = new_segmentation(run, voxel_spacing, object_name, session_id, user_id, "instance") + output.from_numpy(volume) + stats = {"filaments_painted": len(ids), "voxels_painted": int(np.count_nonzero(volume))} + logger.info(f"Painted {len(ids)} filament tubes ({stats['voxels_painted']} voxels) in {run.name}") + return output, stats + + except Exception as e: + logger.error(f"Error painting filaments in {run.name}: {e}") + return None + + +segmentation_from_filaments_lazy_batch = create_lazy_batch_converter( + converter_func=segmentation_from_filaments, + task_description="Painting filaments", +) diff --git a/src/copick_utils/process/__init__.py b/src/copick_utils/process/__init__.py index ed8669a..29e3084 100644 --- a/src/copick_utils/process/__init__.py +++ b/src/copick_utils/process/__init__.py @@ -1,6 +1,7 @@ """Segmentation processing utilities for copick.""" from .connected_components import ( + components_as_instances, extract_individual_components, print_component_stats, separate_components_converter, @@ -8,6 +9,14 @@ separate_connected_components_3d, separate_segmentation_components, ) +from .filament_tracing import ( + Centreline, + TraceParameters, + TraceReport, + paint_tubes, + sample_curve, + trace_centrelines, +) from .rescale import ( rescale_lazy_batch, rescale_segmentation, @@ -27,6 +36,7 @@ from .split_labels import ( split_labels_batch, split_multilabel_segmentation, + split_panoptic_segmentation, ) from .thickness_filter import ( thickness_filter_lazy_batch, @@ -45,6 +55,7 @@ "separate_segmentation_components", "separate_components_converter", "separate_components_lazy_batch", + "components_as_instances", "TubeSkeletonizer3D", "skeletonize_segmentation", "skeletonize_converter", @@ -53,6 +64,13 @@ "rescale_lazy_batch", "split_multilabel_segmentation", "split_labels_batch", + "split_panoptic_segmentation", + "Centreline", + "TraceParameters", + "TraceReport", + "trace_centrelines", + "sample_curve", + "paint_tubes", "SkeletonSplineFitter", "fit_spline_to_skeleton", "fit_spline_to_segmentation", diff --git a/src/copick_utils/process/filament_tracing.py b/src/copick_utils/process/filament_tracing.py new file mode 100644 index 0000000..95dbb1a --- /dev/null +++ b/src/copick_utils/process/filament_tracing.py @@ -0,0 +1,994 @@ +"""Filament tracing: centrelines from segmentations, picks along centrelines, and tube masks around them. + +The functions here work on numpy arrays and do not read or write copick objects; the converters +(``converters/filaments_from_segmentation.py``, ``picks_from_filaments.py``, ``segmentation_from_filaments.py``) and +``fit-spline`` call them. Curves are evaluated with copick's reference implementation (``copick.util.filaments``), so +a curve is sampled the same way copick regenerates its points. + +Conventions (copick ``docs/geometry.md``): + +- Volumes are indexed (z, y, x); voxel ``i`` sits at ``i * voxel_size`` Angstrom. Positions are (x, y, z) Angstrom. +- A filament's points are ordered along it; its picks carry its ID as ``instance_id``, and each pick's rotation has + its +Z axis along the filament (the tangent, in point order). The rotation about the axis (roll) is documented by + ``sample_curve``. + +Tracing a semantic segmentation (``trace_centrelines``): + +1. Connected components (26-connected) smaller than ``min_volume`` are dropped. +2. Each remaining component is skeletonized (Lee) and turned into a skeleton graph (``skeleton_graph.py``); spurs + shorter than ``prune_length`` are pruned, and junctions joined by a bridge shorter than ``junction_merge`` are + merged (two crossing filaments often skeletonize as two forks joined by a short bridge). +3. At every junction, branches are paired by straightest continuation: the pair whose directions deviate least from + a straight line is joined first, as long as the deviation is at most ``max_bend`` degrees. A crossing therefore + gives two filaments, and a Y gives one filament through the junction and one that ends there. +4. Paired branches are chained into filaments. At a free end, the last label radius of skeleton is replaced by the + continuation of the filament's direction up to the edge of the label (thinning shortens a rod's skeleton by about + one radius at each end and often bends its last voxels). +5. Filaments shorter than ``min_length``, shorter than ``min_aspect`` label diameters (blobs), or thinner than + ``min_radius`` (slivers of segmentation noise, much thinner than the real filaments) are rejected. +6. Each filament is fitted with a smoothing B-spline (scipy ``splprep``) whose RMS deviation from the skeleton is + about ``smoothing`` Angstrom. The fit is returned exactly (``tck``), so copick can store it as a ``bspline`` curve. +7. Every label voxel is assigned to the nearest filament of its component, which gives the instance segmentation. + +Tracing an instance segmentation traces each instance on its own and keeps its ID; an instance that branches keeps +its longest chain. +""" + +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Sequence, Tuple + +import numpy as np +from scipy import ndimage +from scipy.interpolate import splev, splprep +from scipy.spatial import cKDTree + +from copick_utils.process.skeleton_graph import ( + Topology, + decompose, + merge_junctions, + node_centre, + prune_spurs, + skeleton_graph, + splice, +) + +#: Version of the tracing algorithm, recorded in each filament's ``metadata["fit"]``. +TRACER_VERSION = 1 + + +@dataclass +class TraceParameters: + """Tracing parameters. Lengths are in Angstrom and volumes in cubic Angstrom; ``None`` derives a value from the + label itself, as noted, so the defaults carry over between objects and voxel sizes. + + Attributes: + min_volume: Drop connected components smaller than this. None: no volume filter. + min_length: Reject filaments shorter than this. None: no length filter. + min_aspect: Reject filaments shorter than this many label diameters (compact blobs). Basis: on dataset 10521 + (easymode microtubule labels, 58 tomograms) every microtubule-thick piece this rejects is under 44 nm, + and every piece of 46 nm or more passes. + min_radius: Reject filaments whose median label radius is below this (thin slivers of segmentation noise). + None: no radius filter. Basis: on dataset 10521, real microtubules have a label radius of 40-76 Å, while + slivers and stubs of segmentation noise have 10-35 Å; the converters default to a third of the object's + tube radius (40 Å there). + prune_length: Prune skeleton spurs shorter than this. None: one label diameter. + junction_merge: Merge junctions joined by a bridge up to this long. None: one label diameter. + max_bend: Largest deviation from a straight line, in degrees, for two branches to continue through a junction. + smoothing: RMS deviation of the fitted spline from the skeleton. None: half a voxel. + extend_ends: Extend free ends along their direction to the edge of the label. + direction_window: Length over which a branch's direction at a junction is measured. None: two label + diameters. + fill_lumen: Before skeletonizing, fill holes up to this radius in every axis-aligned slice of the label. A + label of a tube's wall alone (an empty lumen) otherwise skeletonizes to a mesh over the wall. None: no + filling; the converters default to the object's tube radius. Basis: on dataset 10521, ten tomograms have + stretches of hollow label; without filling, three of them traced no microtubule at all. + """ + + min_volume: Optional[float] = None + min_length: Optional[float] = None + min_aspect: float = 3.0 + min_radius: Optional[float] = None + prune_length: Optional[float] = None + junction_merge: Optional[float] = None + max_bend: float = 45.0 + smoothing: Optional[float] = None + extend_ends: bool = True + direction_window: Optional[float] = None + fill_lumen: Optional[float] = None + + def as_dict(self) -> Dict[str, Any]: + """The parameters as a plain dict (for ``metadata["fit"]``).""" + return dict(self.__dict__) + + +@dataclass +class Centreline: + """A traced filament. + + Attributes: + instance_id: The filament's ID (>= 1). + tck: The fitted spline as scipy's ``(knots, [cx, cy, cz], degree)``, in Angstrom (x, y, z). + smoothing: The ``splprep`` smoothing factor ``s`` of the fit (Angstrom squared). + points: Dense points along the fit, (N, 3) Angstrom (x, y, z), at most one voxel apart. + length: Length along the fit, in Angstrom. + radius: Median label radius along the skeleton (distance to the background), in Angstrom. + score: Fraction of the fitted centreline that lies inside the label. + component: The connected component (or input instance ID) the filament came from. + metadata: Details for ``metadata["fit"]`` (junctions, extended ends, closed, ...). + """ + + instance_id: int + tck: Tuple[np.ndarray, List[np.ndarray], int] + smoothing: float + points: np.ndarray + length: float + radius: Optional[float] + score: float + component: int + metadata: Dict[str, Any] = field(default_factory=dict) + + +@dataclass +class TraceReport: + """Counts from one tracing run.""" + + components: int = 0 + components_removed_volume: int = 0 + filaments: int = 0 + rejected_length: int = 0 + rejected_aspect: int = 0 + rejected_radius: int = 0 + junctions_resolved: int = 0 + total_length: float = 0.0 + dropped_branch_length: float = 0.0 + + def as_dict(self) -> Dict[str, float]: + """The counts as a dict (numeric, so batch runs can sum them).""" + return dict(self.__dict__) + + +# --------------------------------------------------------------------------------------------------------------------- +# Skeleton chains +# --------------------------------------------------------------------------------------------------------------------- + + +def _branch_end(topology: Topology, branch_index: int, at_start: bool) -> np.ndarray: + """The voxel position (z, y, x) where a branch meets its node.""" + branch = topology.branches[branch_index] + return topology.graph.coords[branch.path[0] if at_start else branch.path[-1]].astype(float) + + +def _end_direction(topology: Topology, branch_index: int, at_start: bool, window: float) -> np.ndarray: + """Unit direction (z, y, x) pointing from a junction into a branch, measured over ``window`` voxels of path from + the branch's own end (not the junction's centre, which for a cluster of merged junctions can lie far from it).""" + branch = topology.branches[branch_index] + path = branch.path if at_start else branch.path[::-1] + coords = topology.graph.coords[path].astype(float) + steps = np.linalg.norm(np.diff(coords, axis=0), axis=1) + reach = np.concatenate([[0.0], np.cumsum(steps)]) + far = coords[min(int(np.searchsorted(reach, window)), len(coords) - 1)] + direction = far - coords[0] + norm = np.linalg.norm(direction) + if norm < 1e-9: + direction = far - node_centre(topology, branch.start if at_start else branch.end) + norm = np.linalg.norm(direction) + return direction / norm if norm > 1e-9 else np.zeros(3) + + +def _angle(a: np.ndarray, b: np.ndarray) -> float: + """Angle in degrees between two unit vectors.""" + return float(np.degrees(np.arccos(np.clip(np.dot(a, b), -1.0, 1.0)))) + + +def _pair_junction( + topology: Topology, + node: int, + max_bend: float, + window: float, + gap_tolerance: float = 2.0, +) -> List[Tuple[Tuple[int, bool], Tuple[int, bool]]]: + """Pair the branch ends at a junction by straightest continuation, greedily, within ``max_bend`` degrees. + + Continuing from one branch into another means arriving along the first, crossing the junction from its end to + the second's, and leaving along the second. Where the two ends are more than ``gap_tolerance`` voxels apart (a + cluster of merged junctions), the crossing has to continue straight as well, so a pair that turns back inside + the junction is never chosen. + """ + ends = topology.incident(node) + directions = [_end_direction(topology, b, at_start, window) for b, at_start in ends] + positions = [_branch_end(topology, b, at_start) for b, at_start in ends] + candidates = [] + for i in range(len(ends)): + for j in range(i + 1, len(ends)): + arrive, leave = -directions[i], directions[j] # symmetric: reversing both gives the same bend + bend = _angle(arrive, leave) + gap = positions[j] - positions[i] + distance = np.linalg.norm(gap) + if distance > gap_tolerance: + gap = gap / distance + bend = max(bend, _angle(arrive, gap), _angle(gap, leave)) + candidates.append((bend, i, j)) + pairs, used = [], set() + for bend, i, j in sorted(candidates): + if bend > max_bend: + break + if i in used or j in used: + continue + used.update((i, j)) + pairs.append((ends[i], ends[j])) + return pairs + + +@dataclass +class Chain: + """An ordered skeleton path through one or more branches. + + Attributes: + coords: (N, 3) voxel positions (z, y, x), in order. + free_ends: Whether the chain's start and end are free (endpoints), as opposed to ending at a junction. + closed: Whether the chain closes on itself. + junctions: Number of junctions the chain passes through. + """ + + coords: np.ndarray + free_ends: Tuple[bool, bool] + closed: bool = False + junctions: int = 0 + + +def skeleton_chains( + skeleton: np.ndarray, + prune_length: float, + junction_merge: float, + max_bend: float, + direction_window: float, + junction_trim: float = 0.0, +) -> Tuple[List[Chain], int]: + """Split a skeleton into filament chains, continuing straight through junctions. + + Args: + skeleton: 3D boolean skeleton (one connected component, or several). + prune_length: Prune spurs shorter than this many voxels. + junction_merge: Merge junctions joined by a bridge up to this many voxels. + max_bend: Largest deviation from straight, in degrees, to continue through a junction. + direction_window: Voxels of path over which a branch's direction at a junction is measured. + junction_trim: Drop this many voxels of path where a branch meets a junction. Thinning bends each branch + into the junction's centre; without those voxels, a chain through a junction joins its branches directly, + and a chain that ends at a junction ends straight. + + Returns: + Tuple of (chains, number of junctions resolved). + """ + topology = splice(decompose(skeleton_graph(skeleton))) + if prune_length > 0: + topology = prune_spurs(topology, prune_length) + if junction_merge > 0: + topology = splice(merge_junctions(topology, junction_merge)) + + # Links between branch ends that continue through a junction + link: Dict[Tuple[int, bool], Tuple[int, bool]] = {} + junctions = topology.junctions() + for node in junctions: + for a, b in _pair_junction(topology, node, max_bend, direction_window): + link[a] = b + link[b] = a + + coords = topology.graph.coords + visited = set() + chains: List[Chain] = [] + + def doubles_back(parts: List[np.ndarray], branch: int, at_start: bool) -> bool: + """Whether a branch, beyond the junction it starts from, runs back along the chain so far.""" + if not parts or junction_trim <= 0: + return False + path = coords[topology.branches[branch].path if at_start else topology.branches[branch].path[::-1]] + steps = np.concatenate([[0.0], np.cumsum(np.linalg.norm(np.diff(path, axis=0), axis=1))]) + beyond = path[steps > 4.0 * junction_trim].astype(float) + if len(beyond) == 0: + return False + earlier = np.concatenate(parts) + steps_back = np.concatenate([[0.0], np.cumsum(np.linalg.norm(np.diff(earlier[::-1], axis=0), axis=1))]) + earlier = earlier[::-1][steps_back > 4.0 * junction_trim] + if len(earlier) == 0: + return False + distance, _ = cKDTree(earlier).query(beyond) + return bool(np.any(distance < 2.0 * junction_trim)) + + def follow(start_branch: int, enter_at_start: bool) -> Chain: + parts, n_junctions = [], 0 + branch, at_start = start_branch, enter_at_start + first_end = (branch, at_start) + closed = False + while True: + visited.add(branch) + current = topology.branches[branch] + path = coords[current.path if at_start else current.path[::-1]].astype(float) + entry_node = current.start if at_start else current.end + exit_node = current.end if at_start else current.start + if junction_trim > 0 and topology.node_kind.get(entry_node) == "junction": + path = _trim(path, junction_trim, at_start=True) + if junction_trim > 0 and topology.node_kind.get(exit_node) == "junction": + path = _trim(path, junction_trim, at_start=False) + parts.append(path) + exit_end = (branch, not at_start) + if exit_end not in link: + break + nxt = link[exit_end] + if nxt == first_end or nxt[0] in visited: + closed = nxt == first_end + n_junctions += int(closed) + break + if doubles_back(parts, *nxt): + break + n_junctions += 1 + branch, at_start = nxt + free_start = first_end not in link + free_end = not closed and (branch, not at_start) not in link + start_kind = topology.node_kind.get( + topology.branches[first_end[0]].start if first_end[1] else topology.branches[first_end[0]].end, + ) + last = topology.branches[branch] + end_kind = topology.node_kind.get(last.end if at_start else last.start) + return Chain( + coords=np.concatenate(parts), + free_ends=(free_start and start_kind == "endpoint", free_end and end_kind == "endpoint"), + closed=closed or (len(parts) == 1 and topology.branches[first_end[0]].closed), + junctions=n_junctions, + ) + + # Open chains start at an end that continues nowhere + for b in range(len(topology.branches)): + if b in visited: + continue + for at_start in (True, False): + if (b, at_start) not in link: + chains.append(follow(b, at_start)) + break + # What is left are closed loops of linked branches + for b in range(len(topology.branches)): + if b not in visited: + chains.append(follow(b, True)) + return chains, len(link) // 2 + + +# --------------------------------------------------------------------------------------------------------------------- +# Fitting +# --------------------------------------------------------------------------------------------------------------------- + + +def _dedupe(points: np.ndarray) -> np.ndarray: + """Drop consecutive duplicate points.""" + if len(points) < 2: + return points + keep = np.concatenate([[True], np.linalg.norm(np.diff(points, axis=0), axis=1) > 1e-9]) + return points[keep] + + +def fit_spline(points: np.ndarray, s: float, degree: int = 3) -> Optional[Tuple[np.ndarray, List[np.ndarray], int]]: + """A smoothing B-spline through ordered points (scipy ``splprep``, chord-length parameterisation). + + Args: + points: (N, 3) ordered points. + s: ``splprep`` smoothing factor (sum of squared deviations, in the points' units squared). + degree: Spline degree (lowered for very short inputs). + + Returns: + ``(knots, [cx, cy, cz], degree)`` with clamped knots, or None for fewer than two distinct points. + """ + points = _dedupe(np.asarray(points, dtype=float)) + if len(points) < 2: + return None + k = int(min(degree, len(points) - 1)) + chord = np.concatenate([[0.0], np.cumsum(np.linalg.norm(np.diff(points, axis=0), axis=1))]) + u = chord / chord[-1] + tck, _ = splprep(points.T, u=u, s=float(s), k=k) + knots, coefficients, k = tck + count = len(knots) - int(k) - 1 # scipy may pad the coefficient arrays past the spline's own count + return np.asarray(knots, dtype=float), [np.asarray(c, dtype=float)[:count] for c in coefficients], int(k) + + +def evaluate_spline( + tck: Tuple[np.ndarray, List[np.ndarray], int], + step: float, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Dense points, unit tangents and arc length along a B-spline. + + Args: + tck: ``(knots, [cx, cy, cz], degree)``. + step: Largest spacing of the returned points (the spline's units). + + Returns: + Tuple of (points (N, 3), tangents (N, 3), arc length at each point (N,)). + """ + knots, coefficients, k = tck + u0, u1 = knots[k], knots[-k - 1] + coarse = np.array(splev(np.linspace(u0, u1, 64), (knots, coefficients, k))).T + rough = float(np.sum(np.linalg.norm(np.diff(coarse, axis=0), axis=1))) + n = max(16, int(np.ceil(4 * rough / step)) + 1) + u = np.linspace(u0, u1, n) + points = np.array(splev(u, (knots, coefficients, k))).T + tangents = np.array(splev(u, (knots, coefficients, k), der=1)).T + tangents /= np.maximum(np.linalg.norm(tangents, axis=1, keepdims=True), 1e-12) + arc = np.concatenate([[0.0], np.cumsum(np.linalg.norm(np.diff(points, axis=0), axis=1))]) + return points, tangents, arc + + +def _trim(coords: np.ndarray, length: float, at_start: bool) -> np.ndarray: + """Remove ``length`` voxels of path from one end of a chain, keeping at least half of it.""" + path = coords[::-1] if at_start else coords + if len(path) < 4: + return coords + steps = np.linalg.norm(np.diff(path, axis=0), axis=1) + remaining = np.concatenate([[0.0], np.cumsum(steps[::-1])])[::-1] # path length from each point to the end + keep = max(int(np.searchsorted(-remaining, -length)), len(path) // 2) + keep = min(keep, len(path)) + trimmed = path[:keep] if keep >= 2 else path + return trimmed[::-1] if at_start else trimmed + + +def _extend( + coords: np.ndarray, + mask: np.ndarray, + window: float, + limit: float, + at_start: bool, +) -> Tuple[np.ndarray, bool]: + """Extend a chain end (voxel coords) along its direction while inside ``mask``, by at most ``limit`` voxels.""" + path = coords[::-1] if at_start else coords + if len(path) < 2: + return coords, False + steps = np.linalg.norm(np.diff(path[::-1], axis=0), axis=1) + reach = np.concatenate([[0.0], np.cumsum(steps)]) + back = path[::-1][min(int(np.searchsorted(reach, window)), len(path) - 1)] + direction = path[-1] - back + norm = np.linalg.norm(direction) + if norm < 1e-9: + return coords, False + direction /= norm + added, shape = [], np.array(mask.shape) + for t in np.arange(0.5, limit + 1e-9, 0.5): + point = path[-1] + t * direction + voxel = np.round(point).astype(int) + if np.any(voxel < 0) or np.any(voxel >= shape) or not mask[tuple(voxel)]: + break + added.append(point) + if not added: + return coords, False + added = np.array(added) + return (np.concatenate([added[::-1], coords]) if at_start else np.concatenate([coords, added])), True + + +# --------------------------------------------------------------------------------------------------------------------- +# Tracing +# --------------------------------------------------------------------------------------------------------------------- + + +def fill_lumen(mask: np.ndarray, max_radius: float) -> np.ndarray: + """Fill the holes of a 3D mask that are enclosed within a slice along any axis and no larger than a disc of + ``max_radius`` voxels, such as the lumen of a tube whose wall alone is labelled. Larger enclosed regions (a loop + of several filaments) are left open. + + Args: + mask: 3D boolean mask. + max_radius: Radius, in voxels, of the largest disc-sized hole to fill. + + Returns: + The filled mask (a copy). + """ + max_area = np.pi * max_radius**2 + filled = mask.copy() + for axis in range(3): + for index in range(mask.shape[axis]): + plane = np.take(mask, index, axis=axis) + holes = ndimage.binary_fill_holes(plane) & ~plane + if not holes.any(): + continue + labels, _ = ndimage.label(holes) + small = np.bincount(labels.ravel()) <= max_area + small[0] = False + target = [slice(None)] * 3 + target[axis] = index + filled[tuple(target)] |= small[labels] + return filled + + +def _trace_mask( + mask: np.ndarray, + offset: np.ndarray, + voxel_size: float, + params: TraceParameters, + single: bool, + component: int, + report: TraceReport, +) -> List[Centreline]: + """Trace the filaments of one component (or instance) mask, cropped at ``offset`` (z, y, x).""" + from skimage.morphology import skeletonize + + padded = np.pad(mask, 1) + skeleton = skeletonize(padded) + if not skeleton.any(): + # Thinning removes compact specks and blobs entirely: they have no length. + report.rejected_length += 1 + return [] + edt = ndimage.distance_transform_edt(padded) + on_skeleton = edt[skeleton] + radius = float(np.median(on_skeleton)) if on_skeleton.size else 0.5 # voxels + diameter = 2.0 * radius + + def voxels(value: Optional[float], default: float) -> float: + return default if value is None else float(value) / voxel_size + + chains, resolved = skeleton_chains( + skeleton, + prune_length=voxels(params.prune_length, diameter), + junction_merge=voxels(params.junction_merge, diameter), + max_bend=params.max_bend, + direction_window=voxels(params.direction_window, 2.0 * diameter), + junction_trim=radius, + ) + report.junctions_resolved += resolved + if single and len(chains) > 1: + lengths = [np.sum(np.linalg.norm(np.diff(c.coords, axis=0), axis=1)) for c in chains] + keep = int(np.argmax(lengths)) + report.dropped_branch_length += float((sum(lengths) - lengths[keep]) * voxel_size) + chains = [chains[keep]] + + sigma = voxels(params.smoothing, 0.5) * voxel_size # Angstrom + window = voxels(params.direction_window, 2.0 * diameter) + centrelines = [] + for chain in chains: + coords = chain.coords.astype(float) + extended = [] + if params.extend_ends and not chain.closed: + for at_start, free in zip((True, False), chain.free_ends): + if free: + # Thinning bends the last radius or so of a skeleton (hooks at ends, and where the label meets + # the volume's edge), so drop it and continue the filament's own direction to the label's edge. + coords = _trim(coords, radius, at_start) + coords, did = _extend(coords, padded, window, 2.5 * radius + 1.0, at_start) + extended.append(did) + coords = coords - 1.0 + offset # undo the padding, back to the volume's (z, y, x) + points = coords[:, ::-1] * voxel_size # (x, y, z) Angstrom + length = float(np.sum(np.linalg.norm(np.diff(points, axis=0), axis=1))) if len(points) > 1 else 0.0 + local = edt[tuple(np.clip(np.round(chain.coords).astype(int), 0, np.array(padded.shape) - 1).T)] + chain_radius = float(np.median(local)) * voxel_size if local.size else None + if params.min_length is not None and length < params.min_length: + report.rejected_length += 1 + continue + if params.min_radius is not None and (chain_radius or 0.0) < params.min_radius: + report.rejected_radius += 1 + continue + if chain_radius and length < params.min_aspect * 2.0 * chain_radius: + report.rejected_aspect += 1 + continue + points = _dedupe(points) + s = len(points) * sigma**2 + tck = fit_spline(points, s=s) + if tck is None: + report.rejected_length += 1 + continue + tck, s, hooks = _fit_without_end_hooks(points, tck, sigma, diameter * voxel_size, params.max_bend) + dense, _, arc = evaluate_spline(tck, step=voxel_size) + inside = _inside(dense, padded, offset, voxel_size) + centrelines.append( + Centreline( + instance_id=0, + tck=tck, + smoothing=s, + points=dense, + length=float(arc[-1]), + radius=chain_radius, + score=float(inside), + component=component, + metadata={ + "junctions": chain.junctions, + "ends_extended": int(sum(extended)), + "end_hooks_removed": hooks, + "closed": bool(chain.closed), + }, + ), + ) + return centrelines + + +def _end_bend(tck: Tuple[np.ndarray, List[np.ndarray], int], diameter: float) -> Tuple[float, float]: + """The angle (degrees) at each end between the curve's end tangent and its direction from one to three label + diameters in; a large angle is a hook at that end.""" + points, tangents, arc = evaluate_spline(tck, step=max(diameter / 8.0, 1e-6)) + length = arc[-1] + if length < 4.0 * diameter: + return 0.0, 0.0 + + def direction(a: float, b: float) -> np.ndarray: + i, j = np.searchsorted(arc, a), np.searchsorted(arc, b) + d = points[min(j, len(points) - 1)] - points[min(i, len(points) - 1)] + return d / max(np.linalg.norm(d), 1e-12) + + start = np.degrees(np.arccos(np.clip(np.dot(tangents[0], direction(diameter, 3.0 * diameter)), -1.0, 1.0))) + inner_end = direction(length - 3.0 * diameter, length - diameter) + end = np.degrees(np.arccos(np.clip(np.dot(tangents[-1], inner_end), -1.0, 1.0))) + return float(start), float(end) + + +def _fit_without_end_hooks( + points: np.ndarray, + tck: Tuple[np.ndarray, List[np.ndarray], int], + sigma: float, + diameter: float, + max_bend: float, + rounds: int = 3, +) -> Tuple[Tuple[np.ndarray, List[np.ndarray], int], float, int]: + """Refit without the end points that make the curve hook at an end (thinning artefacts where a filament meets + another structure or the edge of the volume). An end is trimmed by one label diameter each round, at most + ``rounds`` times. + + Returns: + Tuple of (spline, its smoothing factor, number of end trims). + """ + trims = 0 + s = len(points) * sigma**2 + for _ in range(rounds): + start_bend, end_bend = _end_bend(tck, diameter) + if start_bend <= max_bend and end_bend <= max_bend: + break + chord = np.concatenate([[0.0], np.cumsum(np.linalg.norm(np.diff(points, axis=0), axis=1))]) + keep = np.ones(len(points), dtype=bool) + if start_bend > max_bend: + keep &= chord >= diameter + trims += 1 + if end_bend > max_bend: + keep &= chord <= chord[-1] - diameter + trims += 1 + if keep.sum() < 4: + break + points = points[keep] + s = len(points) * sigma**2 + refit = fit_spline(points, s=s) + if refit is None: + break + tck = refit + return tck, s, trims + + +def _inside(points: np.ndarray, padded: np.ndarray, offset: np.ndarray, voxel_size: float) -> float: + """Fraction of points (x, y, z Angstrom) inside a padded mask cropped at ``offset``.""" + zyx = np.round(points[:, ::-1] / voxel_size - offset + 1.0).astype(int) + valid = np.all((zyx >= 0) & (zyx < np.array(padded.shape)), axis=1) + hits = np.zeros(len(zyx), dtype=bool) + hits[valid] = padded[tuple(zyx[valid].T)] + return float(hits.mean()) if len(hits) else 0.0 + + +def _orient(centreline: Centreline, voxel_size: float) -> Centreline: + """Order a filament so that it starts at the end with the smaller (z, y, x) voxel; reverse the fit if needed.""" + start = np.round(centreline.points[0][::-1] / voxel_size).astype(int) + end = np.round(centreline.points[-1][::-1] / voxel_size).astype(int) + if tuple(end) < tuple(start): + knots, coefficients, k = centreline.tck + u0, u1 = knots[0], knots[-1] + centreline.tck = (u0 + u1 - knots[::-1], [c[::-1] for c in coefficients], k) + centreline.points = centreline.points[::-1] + return centreline + + +def trace_centrelines( + volume: np.ndarray, + voxel_size: float, + params: Optional[TraceParameters] = None, + instances: bool = False, +) -> Tuple[List[Centreline], np.ndarray, TraceReport]: + """Trace the filaments of a segmentation. + + Args: + volume: (Z, Y, X) segmentation. Without ``instances``, any non-zero voxel is filament. With ``instances``, each + positive value is one filament (an instance segmentation). + voxel_size: Voxel size in Angstrom. + params: Tracing parameters (defaults: ``TraceParameters()``). + instances: Trace each instance on its own and keep its ID. + + Returns: + Tuple of (centrelines, assignment, report): + - centrelines: The filaments, with IDs 1..K by length (longest first) for a semantic input, or the input's + IDs for an instance input. Points start at the end with the smaller (z, y, x). + - assignment: (Z, Y, X) instance segmentation: each label voxel holds the ID of the nearest filament of its + component; voxels of components without a filament are 0. For an instance input, the kept instances. + With ``fill_lumen``, the filled holes are part of the label. + - report: Counts. + """ + params = params or TraceParameters() + report = TraceReport() + volume = np.asarray(volume) + if params.fill_lumen: + volume = _fill_label(volume, params.fill_lumen / voxel_size, instances) + if instances: + labeled = volume.astype(np.int64, copy=False) + ids = np.unique(labeled[labeled > 0]) + else: + labeled, n = ndimage.label(volume > 0, structure=np.ones((3, 3, 3))) + ids = np.arange(1, n + 1) + counts = np.bincount(labeled.ravel()) if labeled.size else np.zeros(1, dtype=int) + boxes = ndimage.find_objects(labeled) + report.components = int(len(ids)) + + found: List[Centreline] = [] + for value in ids: + value = int(value) + if params.min_volume is not None and counts[value] * voxel_size**3 < params.min_volume: + report.components_removed_volume += 1 + continue + box = boxes[value - 1] + if box is None: + continue + offset = np.array([sl.start for sl in box]) + mask = labeled[box] == value + if instances: + mask, _ = _largest_piece(mask) + traced = _trace_mask(mask, offset, voxel_size, params, single=instances, component=value, report=report) + for centreline in traced: + if instances: + centreline.instance_id = value + found.append(_orient(centreline, voxel_size)) + + if not instances: + found.sort(key=lambda c: -c.length) + for i, centreline in enumerate(found, start=1): + centreline.instance_id = i + report.filaments = len(found) + report.total_length = float(sum(c.length for c in found)) + for centreline in found: + centreline.metadata.update( + { + "method": "splprep", + "degree": int(centreline.tck[2]), + "smoothing": float(centreline.smoothing), + "component": int(centreline.component), + "label_radius": centreline.radius, + "tracer_version": TRACER_VERSION, + "parameters": params.as_dict(), + }, + ) + + if instances: + kept = np.zeros(int(labeled.max()) + 1 if labeled.size else 1, dtype=bool) + kept[[c.instance_id for c in found]] = True + assignment = np.where(kept[labeled], labeled, 0) + else: + assignment = _assign(labeled, found, boxes, voxel_size) + return found, assignment, report + + +def _fill_label(volume: np.ndarray, max_radius: float, instances: bool) -> np.ndarray: + """The label with its lumen-sized holes filled (see ``fill_lumen``): as a whole for a semantic label, and one + instance at a time, into empty voxels only, for an instance segmentation.""" + if not instances: + return fill_lumen(volume > 0, max_radius) + filled = volume.copy() + for value, box in enumerate(ndimage.find_objects(volume.astype(np.int64, copy=False)), start=1): + if box is None: + continue + mask = fill_lumen(volume[box] == value, max_radius) + region = filled[box] + region[mask & (region == 0)] = value + return filled + + +def _largest_piece(mask: np.ndarray) -> Tuple[np.ndarray, int]: + """The largest 26-connected piece of a mask, and the number of pieces.""" + pieces, n = ndimage.label(mask, structure=np.ones((3, 3, 3))) + if n <= 1: + return mask, n + sizes = np.bincount(pieces.ravel()) + sizes[0] = 0 + return pieces == int(np.argmax(sizes)), n + + +def _assign( + labeled: np.ndarray, + centrelines: Sequence[Centreline], + boxes: Sequence[Optional[Tuple[slice, ...]]], + voxel_size: float, +) -> np.ndarray: + """Give each voxel of a component the ID of the nearest filament traced from that component.""" + dtype = np.uint16 if len(centrelines) < 2**16 else np.uint32 + assignment = np.zeros(labeled.shape, dtype=dtype) + by_component: Dict[int, List[Centreline]] = {} + for centreline in centrelines: + by_component.setdefault(centreline.component, []).append(centreline) + for component, members in by_component.items(): + box = boxes[component - 1] + local = labeled[box] == component + offset = np.array([sl.start for sl in box]) + voxels_zyx = np.argwhere(local) + points = np.concatenate([c.points for c in members]) + owner = np.concatenate([np.full(len(c.points), c.instance_id) for c in members]) + _, nearest = cKDTree(points).query((voxels_zyx + offset)[:, ::-1] * voxel_size) + sub = assignment[box] + sub[tuple(voxels_zyx.T)] = owner[nearest] + return assignment + + +# --------------------------------------------------------------------------------------------------------------------- +# Sampling +# --------------------------------------------------------------------------------------------------------------------- + + +def evaluate_curve_with_tangents(curve: Dict[str, Any], step: float) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Dense points, unit tangents and arc length along a copick filament curve. + + Args: + curve: A copick curve's fields: ``kind`` (``bspline``, ``catmull-rom`` or ``linear``), ``control_points`` + (Angstrom), and ``degree``/``knots`` (``bspline``) or ``alpha`` (``catmull-rom``). + step: Largest spacing of the returned points, in Angstrom. + + Returns: + Tuple of (points (N, 3), tangents (N, 3), arc length (N,)). A ``bspline`` has its exact derivative as + tangent; the other kinds take the direction of the curve, evaluated by copick at ``step``. + """ + kind = curve["kind"] + control = np.asarray(curve["control_points"], dtype=float) + if kind == "bspline": + coefficients = [control[:, i] for i in range(3)] + return evaluate_spline((np.asarray(curve["knots"], dtype=float), coefficients, int(curve["degree"])), step) + from copick.util.filaments import evaluate_curve + + points = evaluate_curve(control, step=step, kind=kind, alpha=curve.get("alpha")) + tangents = np.gradient(points, axis=0) + tangents /= np.maximum(np.linalg.norm(tangents, axis=1, keepdims=True), 1e-12) + arc = np.concatenate([[0.0], np.cumsum(np.linalg.norm(np.diff(points, axis=0), axis=1))]) + return points, tangents, arc + + +def _least_aligned_normal(tangent: np.ndarray) -> np.ndarray: + """A unit vector perpendicular to ``tangent``, from the coordinate axis least aligned with it.""" + axis = np.eye(3)[int(np.argmin(np.abs(tangent)))] + normal = axis - np.dot(axis, tangent) * tangent + return normal / np.linalg.norm(normal) + + +def rotation_minimizing_frames(points: np.ndarray, tangents: np.ndarray) -> np.ndarray: + """Rotation-minimizing frames along a curve (double reflection; Wang et al. 2008). + + Args: + points: (N, 3) points along the curve. + tangents: (N, 3) unit tangents. + + Returns: + (N, 3, 3) rotations whose columns are (x, y, z) with z the tangent. The first frame's x axis is the coordinate + axis least aligned with the first tangent, made perpendicular to it; later frames turn as little as possible. + """ + n = len(points) + frames = np.zeros((n, 3, 3)) + x = _least_aligned_normal(tangents[0]) + for i in range(n): + if i > 0: + v1 = points[i] - points[i - 1] + c1 = np.dot(v1, v1) + if c1 > 1e-18: + x_l = x - (2.0 / c1) * np.dot(v1, x) * v1 + t_l = tangents[i - 1] - (2.0 / c1) * np.dot(v1, tangents[i - 1]) * v1 + v2 = tangents[i] - t_l + c2 = np.dot(v2, v2) + x = x_l - (2.0 / c2) * np.dot(v2, x_l) * v2 if c2 > 1e-18 else x_l + x = x - np.dot(x, tangents[i]) * tangents[i] + norm = np.linalg.norm(x) + x = x / norm if norm > 1e-12 else _least_aligned_normal(tangents[i]) + z = tangents[i] + y = np.cross(z, x) + frames[i] = np.column_stack([x, y, z]) + return frames + + +def sample_curve( + curve: Dict[str, Any], + spacing: float, + anchor: str = "center", + roll: str = "parallel", + seed: Optional[int] = None, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Picks along a filament curve at a fixed spacing, with +Z along the filament. + + Args: + curve: A copick curve's fields (see ``evaluate_curve_with_tangents``). + spacing: Distance between consecutive picks along the curve, in Angstrom. + anchor: ``center`` splits the length left over after the last full step evenly between the two ends; + ``start`` places the first pick at the start of the curve. + roll: ``parallel``: the frames are rotation-minimizing along the curve, starting from the coordinate axis + least aligned with the first tangent (no frame flips, in any direction). ``random``: each pick gets a + uniformly random rotation about the axis. + seed: Random seed for ``roll="random"``. + + Returns: + Tuple of (positions (M, 3) Angstrom, rotations (M, 3, 3), distance along the curve (M,)). A curve shorter + than ``spacing`` gives one pick at its middle. + """ + if spacing <= 0: + raise ValueError("spacing must be positive") + if anchor not in ("center", "start"): + raise ValueError("anchor must be 'center' or 'start'") + if roll not in ("parallel", "random"): + raise ValueError("roll must be 'parallel' or 'random'") + fine = min(float(spacing), float(curve.get("step") or spacing)) / 4.0 + points, tangents, arc = evaluate_curve_with_tangents(curve, fine) + length = float(arc[-1]) + if length < spacing: + along = np.array([length / 2.0]) + else: + count = int(np.floor(length / spacing + 1e-9)) + 1 + margin = (length - (count - 1) * spacing) / 2.0 if anchor == "center" else 0.0 + along = margin + spacing * np.arange(count) + + positions = np.column_stack([np.interp(along, arc, points[:, i]) for i in range(3)]) + sample_tangents = np.column_stack([np.interp(along, arc, tangents[:, i]) for i in range(3)]) + sample_tangents /= np.maximum(np.linalg.norm(sample_tangents, axis=1, keepdims=True), 1e-12) + + frames = rotation_minimizing_frames(points, tangents) + nearest = np.clip(np.searchsorted(arc, along), 0, len(arc) - 1) + rotations = np.zeros((len(along), 3, 3)) + for i, (k, z) in enumerate(zip(nearest, sample_tangents)): + x = frames[k][:, 0] - np.dot(frames[k][:, 0], z) * z + x /= np.linalg.norm(x) + rotations[i] = np.column_stack([x, np.cross(z, x), z]) + if roll == "random": + angles = np.random.default_rng(seed).uniform(0.0, 2.0 * np.pi, len(along)) + for i, angle in enumerate(angles): + c, s = np.cos(angle), np.sin(angle) + rotations[i] = rotations[i] @ np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]]) + return positions, rotations, along + + +# --------------------------------------------------------------------------------------------------------------------- +# Painting +# --------------------------------------------------------------------------------------------------------------------- + + +def paint_tubes( + polylines: Sequence[np.ndarray], + ids: Sequence[int], + radii: Sequence[float], + shape: Tuple[int, int, int], + voxel_size: float, +) -> np.ndarray: + """An instance segmentation of tubes around filament centrelines. + + Args: + polylines: (N, 3) ordered points per filament, in Angstrom (x, y, z), at most a voxel apart. + ids: Filament IDs (>= 1). + radii: Tube radius per filament, in Angstrom. + shape: (Z, Y, X) of the output. + voxel_size: Voxel size in Angstrom. + + Returns: + (Z, Y, X) array: voxels within a tube hold its filament's ID; where tubes overlap, the nearest centreline wins. + """ + dtype = np.uint16 if (max(ids) if len(ids) else 0) < 2**16 else np.uint32 + out = np.zeros(shape, dtype=dtype) + trees: Dict[int, cKDTree] = {} + shape_arr = np.array(shape) + for polyline, fid, radius in zip(polylines, ids, radii): + polyline = np.asarray(polyline, dtype=float) + if len(polyline) == 0 or radius <= 0: + continue + tree = trees[int(fid)] = cKDTree(polyline) + zyx = polyline[:, ::-1] / voxel_size + r = radius / voxel_size + low = np.maximum(np.floor(zyx.min(axis=0) - r).astype(int), 0) + high = np.minimum(np.ceil(zyx.max(axis=0) + r).astype(int) + 1, shape_arr) + if np.any(high <= low): + continue + grid = np.stack(np.meshgrid(*[np.arange(a, b) for a, b in zip(low, high)], indexing="ij"), axis=-1) + grid = grid.reshape(-1, 3) + distance, _ = tree.query(grid[:, ::-1] * voxel_size, distance_upper_bound=radius) + inside = np.isfinite(distance) + cells, distance = grid[inside], distance[inside] + current = out[tuple(cells.T)] + conflict = current > 0 + if np.any(conflict): + other = np.full(conflict.sum(), np.inf) + for owner in np.unique(current[conflict]): + sel = current[conflict] == owner + other[sel], _ = trees[int(owner)].query(cells[conflict][sel][:, ::-1] * voxel_size) + take = np.ones(len(cells), dtype=bool) + take[np.flatnonzero(conflict)] = distance[conflict] < other + cells = cells[take] + out[tuple(cells.T)] = fid + return out diff --git a/src/copick_utils/process/skeleton_graph.py b/src/copick_utils/process/skeleton_graph.py index dc24999..14b80d4 100644 --- a/src/copick_utils/process/skeleton_graph.py +++ b/src/copick_utils/process/skeleton_graph.py @@ -337,7 +337,8 @@ def prune_spurs(topology: Topology, min_length: float) -> Topology: def merge_junctions(topology: Topology, max_length: float) -> Topology: """Contract branches shorter than ``max_length`` voxels that join two different junctions, so two nearby - junctions (two filaments crossing, seen as two forks joined by a short bridge) become one. + junctions (two filaments crossing, seen as two forks joined by a short bridge) become one. Short branches that + then start and end at the merged junction (bubbles of the skeleton, where the label is thick) are absorbed too. Args: topology: The topology (modified in place and returned). @@ -377,6 +378,9 @@ def find(node): if b in bridge_set: continue branch.start, branch.end = find(branch.start), find(branch.end) + if branch.closed and topology.node_kind.get(branch.start) == "junction" and branch.length <= max_length: + merged_voxels[branch.start].append(branch.path) # a bubble inside the merged junction + continue kept.append(branch) topology.branches = kept topology.node_kind = {find(node): kind for node, kind in topology.node_kind.items()} diff --git a/src/copick_utils/process/skeletonize.py b/src/copick_utils/process/skeletonize.py index 1a7b707..6e89552 100644 --- a/src/copick_utils/process/skeletonize.py +++ b/src/copick_utils/process/skeletonize.py @@ -9,6 +9,7 @@ from skimage.morphology import remove_small_objects, skeletonize from copick_utils.converters.lazy_converter import create_lazy_batch_converter +from copick_utils.process.skeleton_graph import pruned_skeleton if TYPE_CHECKING: from copick.models import CopickRun, CopickSegmentation @@ -16,6 +17,15 @@ logger = get_logger(__name__) +def _remove_small(volume: np.ndarray, min_size: int, connectivity: int = 1) -> np.ndarray: + """Remove connected objects with fewer than ``min_size`` voxels (scikit-image's keyword changed in 0.26).""" + import inspect + + if "max_size" in inspect.signature(remove_small_objects).parameters: + return remove_small_objects(volume, max_size=max(int(min_size) - 1, 0), connectivity=connectivity) + return remove_small_objects(volume, min_size=min_size, connectivity=connectivity) + + class TubeSkeletonizer3D: """3D tube skeletonization class based on scikit-image.""" @@ -42,7 +52,7 @@ def preprocess_volume(self, remove_noise: bool = True, min_object_size: int = 10 min_object_size: Minimum size of objects to keep """ if remove_noise and np.any(self.original_volume): - self.original_volume = remove_small_objects(self.original_volume, min_size=min_object_size) + self.original_volume = _remove_small(self.original_volume, min_object_size) def skeletonize(self, method: str = "skimage"): """ @@ -52,7 +62,7 @@ def skeletonize(self, method: str = "skimage"): method: Method to use ('skimage', 'distance_transform') """ if not np.any(self.original_volume): - print("Warning: Volume is empty, creating empty skeleton") + logger.warning("Volume is empty, creating empty skeleton") self.skeleton = np.zeros_like(self.original_volume, dtype=bool) self.skeleton_coords = np.array([]).reshape(0, 3) return @@ -78,24 +88,31 @@ def skeletonize(self, method: str = "skimage"): # Get skeleton coordinates self.skeleton_coords = np.array(np.where(self.skeleton)).T - def post_process_skeleton(self, remove_short_branches: bool = True, min_branch_length: int = 5): + def post_process_skeleton( + self, + remove_short_branches: bool = True, + min_branch_length: int = 5, + prune_length: Optional[float] = None, + ): """ Post-process the skeleton to remove artifacts. Args: - remove_short_branches: Whether to remove short branches - min_branch_length: Minimum length of branches to keep + remove_short_branches: Whether to remove small skeleton pieces. Despite the name, this removes whole + connected skeleton pieces with fewer than ``min_branch_length`` voxels; it does not shorten side + branches of a larger skeleton (``prune_length`` does). + min_branch_length: Minimum number of voxels of a skeleton piece to keep + prune_length: If set, prune side branches (spurs) shorter than this many voxels """ + if prune_length and len(self.skeleton_coords) > 0: + self.skeleton = pruned_skeleton(self.skeleton, prune_length) + self.skeleton_coords = np.array(np.where(self.skeleton)).T if remove_short_branches and len(self.skeleton_coords) > 0: # Remove small (spur) objects from the skeleton. Use full connectivity so a thin, # diagonally-stepping centerline stays a single connected object — with the default # face-connectivity it fragments into tiny pieces that are then all removed (newer # scikit-image removes objects <= min_size), emptying the skeleton entirely. - cleaned_skeleton = remove_small_objects( - self.skeleton, - min_size=min_branch_length, - connectivity=self.skeleton.ndim, - ) + cleaned_skeleton = _remove_small(self.skeleton, min_branch_length, connectivity=self.skeleton.ndim) self.skeleton = cleaned_skeleton self.skeleton_coords = np.array(np.where(self.skeleton)).T @@ -129,6 +146,8 @@ def skeletonize_segmentation( min_branch_length: int = 5, output_session_id: Optional[str] = None, output_user_id: str = "skel", + output_name: Optional[str] = None, + prune_length: Optional[float] = None, ) -> Optional["CopickSegmentation"]: """ Skeletonize a segmentation volume. @@ -142,6 +161,8 @@ def skeletonize_segmentation( min_branch_length: Minimum length of branches to keep output_session_id: Session ID for output segmentation (default: same as input) output_user_id: User ID for output segmentation + output_name: Name of the output segmentation (default: the input's name) + prune_length: If set, prune skeleton side branches shorter than this many voxels Returns: Created skeleton segmentation or None if failed @@ -149,18 +170,18 @@ def skeletonize_segmentation( # Get the segmentation volume volume = segmentation.numpy() if volume is None: - print(f"Error: Could not load segmentation data for {segmentation.run.name}") + logger.error(f"Could not load segmentation data for {segmentation.run.name}") return None run = segmentation.run voxel_size = segmentation.voxel_size - name = segmentation.name + name = output_name or segmentation.name # Use input session_id if no output session_id specified if output_session_id is None: output_session_id = segmentation.session_id - print(f"Skeletonizing segmentation {segmentation.session_id} in run {run.name}") + logger.info(f"Skeletonizing segmentation {segmentation.session_id} in run {run.name}") # Initialize skeletonizer skeletonizer = TubeSkeletonizer3D() @@ -175,11 +196,15 @@ def skeletonize_segmentation( skeletonizer.skeletonize(method=method) # Post-process - skeletonizer.post_process_skeleton(remove_short_branches=remove_short_branches, min_branch_length=min_branch_length) + skeletonizer.post_process_skeleton( + remove_short_branches=remove_short_branches, + min_branch_length=min_branch_length, + prune_length=prune_length, + ) # Get properties properties = skeletonizer.get_skeleton_properties() - print(f"Skeleton properties: {properties['n_voxels']} voxels") + logger.info(f"Skeleton properties: {properties['n_voxels']} voxels") # Create output segmentation try: @@ -195,11 +220,11 @@ def skeletonize_segmentation( # Store the skeleton volume output_seg.from_numpy(skeletonizer.skeleton.astype(np.uint8)) - print(f"Created skeleton segmentation with session_id: {output_session_id}") + logger.info(f"Created skeleton segmentation with session_id: {output_session_id}") return output_seg except Exception as e: - print(f"Error creating skeleton segmentation: {e}") + logger.error(f"Error creating skeleton segmentation: {e}") return None @@ -214,6 +239,7 @@ def skeletonize_converter( min_object_size: int = 50, remove_short_branches: bool = True, min_branch_length: int = 5, + prune_length: Optional[float] = None, **kwargs, ) -> Optional[Tuple["CopickSegmentation", Dict[str, int]]]: """ @@ -233,6 +259,7 @@ def skeletonize_converter( min_object_size: Minimum size of objects to keep during preprocessing. remove_short_branches: Whether to remove short branches from skeleton. min_branch_length: Minimum length of branches to keep. + prune_length: If set, prune skeleton side branches shorter than this, in angstroms. **kwargs: Additional keyword arguments from lazy converter (ignored). Returns: @@ -246,6 +273,8 @@ def skeletonize_converter( remove_short_branches=remove_short_branches, min_branch_length=min_branch_length, output_session_id=session_id, + output_name=object_name, + prune_length=None if not prune_length else prune_length / float(segmentation.voxel_size), output_user_id=user_id, ) diff --git a/src/copick_utils/process/spline_fitting.py b/src/copick_utils/process/spline_fitting.py index 2ea3dde..210878d 100644 --- a/src/copick_utils/process/spline_fitting.py +++ b/src/copick_utils/process/spline_fitting.py @@ -1,320 +1,250 @@ -"""3D spline fitting to skeleton volumes for pick generation with orientations.""" +"""3D spline fitting to skeleton volumes for pick generation with orientations. -from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple +``fit-spline`` fits a smoothing spline to every filament of a skeleton (or segmentation) volume and samples picks +along each one, in voxel units. It is built on the filament tracer (``filament_tracing.py``): the skeleton is split +into chains between ends and junctions, continuing straight through junctions, and each chain gets its own spline. +For new work, ``copick convert seg2fil`` and ``copick convert fil2picks`` separate tracing from sampling and work in +angstroms. +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple -import networkx as nx import numpy as np from copick.util.log import get_logger -from scipy.spatial.distance import pdist, squareform +from scipy.interpolate import splev from copick_utils.converters.lazy_converter import create_lazy_batch_converter +from copick_utils.process.filament_tracing import ( + evaluate_spline, + fit_spline, + rotation_minimizing_frames, + skeleton_chains, +) if TYPE_CHECKING: from copick.models import CopickPicks, CopickRun, CopickSegmentation logger = get_logger(__name__) +#: Chain-splitting settings for skeleton input, in voxels (the input has no thickness to derive them from). +_PRUNE_VOXELS = 3.0 +_MERGE_VOXELS = 3.0 +_MAX_BEND = 45.0 +_WINDOW_VOXELS = 6.0 + class SkeletonSplineFitter: - """3D spline fitting to skeleton coordinates with point sampling and orientation computation.""" + """3D spline fitting to skeleton coordinates with point sampling and orientation computation. + + Coordinates are voxel indices: skeleton coordinates in (z, y, x) order, as ``np.argwhere`` returns them, and + spline points in (x, y, z) order. + """ def __init__(self): + self.spline_tck = None + self.smoothing = None self.skeleton_coords = None - self.ordered_path = None - self.spline_functions = None - self.regularized_points = None + self.ordered_coords = None + self.sampled_points = None self.t_sampled = None + self.tangents = None + self._dense = None def extract_skeleton_coordinates(self, binary_volume: np.ndarray) -> np.ndarray: - """Extract skeleton coordinates from binary volume.""" - self.skeleton_coords = np.array(np.where(binary_volume)).T + """Extract (z, y, x) coordinates of the skeleton voxels.""" + self.skeleton_coords = np.argwhere(binary_volume) return self.skeleton_coords def order_skeleton_points_longest_path(self, coords: np.ndarray, connectivity_radius: float = 2.0) -> np.ndarray: - """Order skeleton points by finding the longest path through the skeleton.""" - if len(coords) <= 2: - self.ordered_path = coords - return coords - - # Build adjacency matrix - distances = squareform(pdist(coords)) - adjacency = (distances <= connectivity_radius) & (distances > 0) - - # Create NetworkX graph - G = nx.from_numpy_array(adjacency) - - # Find endpoints (degree 1 nodes) - endpoints = [node for node, degree in G.degree() if degree == 1] - - if len(endpoints) < 2: - # If no clear endpoints, use the two points that are farthest apart - max_dist_idx = np.unravel_index(np.argmax(distances), distances.shape) - endpoints = [max_dist_idx[0], max_dist_idx[1]] - - # Find the longest path between any two endpoints - longest_path = [] - max_length = 0 - - for i, start in enumerate(endpoints): - for end in endpoints[i + 1 :]: - try: - path = nx.shortest_path(G, start, end) - if len(path) > max_length: - max_length = len(path) - longest_path = path - except nx.NetworkXNoPath: - continue - - # If no path found, try all pairs of nodes - if not longest_path: - for i in range(len(coords)): - for j in range(i + 1, len(coords)): - try: - path = nx.shortest_path(G, i, j) - if len(path) > max_length: - max_length = len(path) - longest_path = path - except nx.NetworkXNoPath: - continue - - # Return ordered coordinates - if longest_path: - self.ordered_path = coords[longest_path] - else: - # Fallback: order by distance from first point - self.ordered_path = self._order_by_nearest_neighbor(coords) + """Order skeleton points along the longest filament chain. - return self.ordered_path + The skeleton is split into chains between ends and junctions, continuing straight through junctions, and the + longest chain is returned in order. - def _order_by_nearest_neighbor(self, coords: np.ndarray) -> np.ndarray: - """Fallback method to order points by nearest neighbor traversal.""" - if len(coords) == 0: - return coords - - ordered = [0] # Start with first point - remaining = list(range(1, len(coords))) - - while remaining: - current_point = coords[ordered[-1]] - distances = [np.linalg.norm(coords[i] - current_point) for i in remaining] - next_idx = remaining[np.argmin(distances)] - ordered.append(next_idx) - remaining.remove(next_idx) + Args: + coords: (N, 3) skeleton voxel coordinates (z, y, x). + connectivity_radius: Unused; skeleton voxels are connected to their 26 neighbours. - return coords[ordered] + Returns: + (M, 3) ordered coordinates (z, y, x) of the longest chain. + """ + chains = _chains_from_coords(coords) + if not chains: + return np.zeros((0, 3)) + self.ordered_coords = max(chains, key=_chain_length) + return self.ordered_coords def fit_regularized_spline( self, coords: np.ndarray, smoothing_factor: Optional[float] = None, degree: int = 3, - ) -> Tuple[np.ndarray, Dict[str, Any]]: - """Fit regularized 3D parametric spline using scipy.interpolate.splprep.""" - if len(coords) < degree + 1: - raise ValueError(f"Need at least {degree + 1} points for degree {degree} spline") - - # Use splprep for parametric 3D spline fitting - from scipy.interpolate import splprep + ) -> Any: + """Fit a smoothing spline to ordered (z, y, x) coordinates. - # Determine spline degree - k = min(degree, len(coords) - 1) - if k <= 0: - k = 1 - - # coords should be transposed for splprep: [x_coords, y_coords, z_coords] - tck, u = splprep([coords[:, 2], coords[:, 1], coords[:, 0]], s=smoothing_factor, k=k) - print(f"Successfully fitted parametric spline with degree {k}, smoothing {smoothing_factor}") - - self.spline_functions = {"tck": tck, "u_original": u, "degree": k} + Args: + coords: (N, 3) ordered coordinates (z, y, x). + smoothing_factor: scipy ``splprep`` smoothing factor ``s`` in voxels squared; None uses scipy's default + ``m - sqrt(2 m)`` for ``m`` points. + degree: Spline degree (1-5). - return u, self.spline_functions + Returns: + The spline as scipy's ``(knots, [cx, cy, cz], degree)`` in (x, y, z) voxel coordinates. + """ + points = np.asarray(coords, dtype=float)[:, ::-1] + m = len(points) + self.smoothing = float(m - np.sqrt(2 * m)) if smoothing_factor is None else float(smoothing_factor) + self.spline_tck = fit_spline(points, s=self.smoothing, degree=degree) + if self.spline_tck is None: + raise ValueError("Not enough distinct points to fit a spline") + self._dense = None + return self.spline_tck def sample_points_along_spline(self, spacing_distance: float) -> np.ndarray: - """Sample points along the spline at regular intervals using arc-length parameterization.""" - if self.spline_functions is None: - raise ValueError("Spline not fitted yet. Call fit_regularized_spline first.") - - # Get the spline representation - from scipy.interpolate import splev - - tck = self.spline_functions["tck"] - - # Estimate total length by sampling densely along the parameter - u_dense = np.linspace(0, 1, 1000) - points_dense = np.column_stack(splev(u_dense, tck)) - - # Calculate cumulative arc lengths - distances = np.zeros(len(points_dense)) - distances[1:] = np.cumsum(np.linalg.norm(np.diff(points_dense, axis=0), axis=1)) - - total_length = distances[-1] - if total_length == 0: - # Fallback: just return the first point - self.regularized_points = np.array([points_dense[0]]) - self.t_sampled = np.array([0.0]) - return self.regularized_points - - # Calculate number of points needed for desired spacing - n_points = max(2, int(np.ceil(total_length / spacing_distance)) + 1) + """Sample points along the spline at exactly ``spacing_distance`` voxels, from its start. - # Sample at regular arc length intervals - target_distances = np.linspace(0, total_length, n_points) - u_sampled = np.interp(target_distances, distances, u_dense) - - # Evaluate spline at sampled parameter values - sampled_points = np.column_stack(splev(u_sampled, tck)) + Args: + spacing_distance: Distance between consecutive points along the spline, in voxels. - self.regularized_points = sampled_points - self.t_sampled = u_sampled # Store parameter values for transform calculation - return sampled_points + Returns: + (N, 3) sampled points (x, y, z) in voxel coordinates. + """ + if self.spline_tck is None: + raise ValueError("Must fit spline first") + if spacing_distance <= 0: + raise ValueError("spacing_distance must be positive") + points, tangents, arc = self._dense_curve() + along = np.arange(0.0, arc[-1] + 1e-9, spacing_distance) + self.t_sampled = along + self.sampled_points = np.column_stack([np.interp(along, arc, points[:, i]) for i in range(3)]) + sampled_tangents = np.column_stack([np.interp(along, arc, tangents[:, i]) for i in range(3)]) + self.tangents = sampled_tangents / np.maximum(np.linalg.norm(sampled_tangents, axis=1, keepdims=True), 1e-12) + return self.sampled_points def compute_transforms(self) -> np.ndarray: - """ - Compute 4x4 transformation matrices for each sampled point. - Follows the ArtiaX pattern: z_align(last_pos, curr_pos).zero_translation().inverse() - The rotation from particle i-1 to particle i is applied to particle i-1. + """Compute 4x4 transformation matrices for each sampled point. + + The rotation's +Z axis is the spline's tangent (its exact derivative) at the point, in the direction of the + sampling. The rotation about the axis is rotation-minimizing along the spline, starting from the coordinate + axis least aligned with the first tangent, so it never flips. The translation is 0. Returns: np.ndarray: [N, 4, 4] array of transformation matrices """ - if self.regularized_points is None: + if self.sampled_points is None: raise ValueError("Must sample points first before computing transforms") - - n_points = len(self.regularized_points) - transforms = np.zeros((n_points, 4, 4)) - - print(f"Computing transforms for {n_points} points using ArtiaX z_align pattern...") - - # Initialize all transforms as identity - for i in range(n_points): - transforms[i] = np.eye(4) - - # ArtiaX pattern: for each particle i, compute rotation from i-1 to i, apply to i-1 - for i in range(1, n_points): - curr_pos = self.regularized_points[i] - last_pos = self.regularized_points[i - 1] - - # z_align(last_pos, curr_pos).zero_translation().inverse() - rotation_matrix = self._z_align_inverse(last_pos, curr_pos) - - # Apply rotation to the PREVIOUS particle (i-1) - transforms[i - 1][:3, :3] = rotation_matrix - - # Handle the last particle - use the same rotation as the previous particle - if n_points > 1: - transforms[n_points - 1][:3, :3] = transforms[n_points - 2][:3, :3] - + points, tangents, arc = self._dense_curve() + frames = rotation_minimizing_frames(points, tangents) + nearest = np.clip(np.searchsorted(arc, self.t_sampled), 0, len(arc) - 1) + transforms = np.tile(np.eye(4), (len(self.sampled_points), 1, 1)) + for i, (k, z) in enumerate(zip(nearest, self.tangents)): + x = frames[k][:, 0] - np.dot(frames[k][:, 0], z) * z + x /= np.linalg.norm(x) + transforms[i, :3, :3] = np.column_stack([x, np.cross(z, x), z]) return transforms - def _z_align_inverse(self, pt1: np.ndarray, pt2: np.ndarray) -> np.ndarray: - """ - Create the inverse of the z_align transformation. - This rotates the z-axis to align with the pt1->pt2 direction. - Based on the z_align algorithm but returns the inverse matrix. - - Args: - pt1: Two 3D points defining the direction vector - pt2: Two 3D points defining the direction vector - - Returns: - np.ndarray: 3x3 rotation matrix (inverse of z_align) - """ - a, b, c = pt2 - pt1 - l = a * a + c * c # noqa - d = l + b * b - epsilon = 1e-10 - - if abs(d) < epsilon: - # Fallback to identity matrix - return np.eye(3) - - l = np.sqrt(l) # noqa - d = np.sqrt(d) - # Create the z_align rotation matrix - xf = np.zeros((3, 3), dtype=np.float64) - xf[1][1] = l / d - - if abs(l) < epsilon: - xf[0][0] = 1.0 - xf[2][1] = -b / d - else: - xf[0][0] = c / l - xf[2][0] = -a / l - xf[0][1] = -(a * b) / (l * d) - xf[2][1] = -(b * c) / (l * d) - - xf[0][2] = a / d - xf[1][2] = b / d - xf[2][2] = c / d - - return xf - def get_spline_properties(self) -> Dict[str, Any]: - """Get properties of the fitted spline.""" - if self.regularized_points is None: + """Properties of the fitted spline: length, sampled point count and spacing, and curvature (1/voxel).""" + if self.spline_tck is None: return {} - - # Calculate total length - distances = np.linalg.norm(np.diff(self.regularized_points, axis=0), axis=1) - total_length = np.sum(distances) - - # Calculate curvature at sampled points - curvatures = [] - if len(self.regularized_points) >= 3: - for i in range(1, len(self.regularized_points) - 1): - p1, p2, p3 = self.regularized_points[i - 1 : i + 2] - v1 = p2 - p1 - v2 = p3 - p2 - # Approximate curvature - cross_prod = np.linalg.norm(np.cross(v1, v2)) - if np.linalg.norm(v1) > 0 and np.linalg.norm(v2) > 0: - curvature = cross_prod / (np.linalg.norm(v1) * np.linalg.norm(v2)) - curvatures.append(curvature) - + points, _, arc = self._dense_curve() + knots, coefficients, k = self.spline_tck + u = np.linspace(knots[k], knots[-k - 1], max(64, len(points))) + d1 = np.array(splev(u, self.spline_tck, der=1)).T + d2 = np.array(splev(u, self.spline_tck, der=2)).T if k >= 2 else np.zeros_like(d1) + speed = np.maximum(np.linalg.norm(d1, axis=1), 1e-12) + curvature = np.linalg.norm(np.cross(d1, d2), axis=1) / speed**3 + n = len(self.sampled_points) if self.sampled_points is not None else 0 return { - "n_points": len(self.regularized_points), - "total_length": total_length, - "average_spacing": ( - total_length / (len(self.regularized_points) - 1) if len(self.regularized_points) > 1 else 0 - ), - "mean_curvature": np.mean(curvatures) if curvatures else 0, - "max_curvature": np.max(curvatures) if curvatures else 0, - "curvatures": curvatures, + "total_length": float(arc[-1]), + "n_sampled_points": n, + "average_spacing": float(arc[-1] / (n - 1)) if n > 1 else 0.0, + "max_curvature": float(curvature.max()) if len(curvature) else 0.0, + "mean_curvature": float(curvature.mean()) if len(curvature) else 0.0, + "smoothing_factor": self.smoothing, } + def max_turn(self) -> float: + """The largest sine of the turning angle between consecutive sampled segments (the --curvature-threshold + quantity; it depends on the sampling spacing).""" + if self.sampled_points is None or len(self.sampled_points) < 3: + return 0.0 + v = np.diff(self.sampled_points, axis=0) + v1, v2 = v[:-1], v[1:] + cross = np.linalg.norm(np.cross(v1, v2), axis=1) + norms = np.maximum(np.linalg.norm(v1, axis=1) * np.linalg.norm(v2, axis=1), 1e-12) + return float(np.max(cross / norms)) + def detect_high_curvature_outliers(self, coords: np.ndarray, curvature_threshold: float = 0.2) -> np.ndarray: - """Detect points that contribute to high curvature.""" - if len(coords) < 4: - return coords + """Deprecated: returns ``coords`` unchanged. - # Calculate curvature at each point - curvatures = [] - for i in range(1, len(coords) - 1): - p1, p2, p3 = coords[i - 1 : i + 2] - v1 = p2 - p1 - v2 = p3 - p2 - cross_prod = np.linalg.norm(np.cross(v1, v2)) - if np.linalg.norm(v1) > 0 and np.linalg.norm(v2) > 0: - curvature = cross_prod / (np.linalg.norm(v1) * np.linalg.norm(v2)) - curvatures.append(curvature) - else: - curvatures.append(0) + Earlier versions deleted skeleton points around sharp bends, which also removed real curves. Sharp bends are + now handled by raising the spline's smoothing (``fit_spline_to_skeleton``). + """ + logger.warning("detect_high_curvature_outliers is deprecated and no longer removes points") + return coords - # Find points with high curvature - curvatures = np.array(curvatures) - high_curvature_mask = curvatures > curvature_threshold + def _dense_curve(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + if self._dense is None: + self._dense = evaluate_spline(self.spline_tck, step=0.25) + return self._dense - if not np.any(high_curvature_mask): - return coords - # Remove points contributing to high curvature (keep first and last) - points_to_keep = [True] + [not high_curvature_mask[i] for i in range(len(high_curvature_mask))] + [True] +def _chain_length(coords: np.ndarray) -> float: + return float(np.sum(np.linalg.norm(np.diff(coords, axis=0), axis=1))) if len(coords) > 1 else 0.0 - filtered_coords = coords[points_to_keep] - removed_count = len(coords) - len(filtered_coords) - print(f"Removed {removed_count} outlier points with high curvature") - return filtered_coords +def _chains_from_coords(coords: np.ndarray) -> List[np.ndarray]: + """Ordered (z, y, x) chains of a skeleton given as voxel coordinates.""" + coords = np.asarray(coords, dtype=int) + if len(coords) == 0: + return [] + low = coords.min(axis=0) + volume = np.zeros(tuple(coords.max(axis=0) - low + 1), dtype=bool) + volume[tuple((coords - low).T)] = True + return [chain.coords + low for chain in _skeleton_chains(volume)] + + +def _skeleton_chains(volume: np.ndarray): + from skimage.morphology import skeletonize + + skeleton = skeletonize(np.pad(np.asarray(volume, dtype=bool), 1))[1:-1, 1:-1, 1:-1] + if not skeleton.any(): + skeleton = np.asarray(volume, dtype=bool) # already as thin as it gets + chains, _ = skeleton_chains( + skeleton, + prune_length=_PRUNE_VOXELS, + junction_merge=_MERGE_VOXELS, + max_bend=_MAX_BEND, + direction_window=_WINDOW_VOXELS, + ) + return chains + + +def _fit_chain( + coords: np.ndarray, + spacing_distance: float, + smoothing_factor: Optional[float], + degree: int, + compute_transforms: bool, + curvature_threshold: float, + max_iterations: int, +) -> Tuple[np.ndarray, Optional[np.ndarray], SkeletonSplineFitter, Dict[str, Any]]: + """Fit one ordered chain; if the sampled curve turns more sharply than ``curvature_threshold``, raise the + smoothing (doubling it, up to ``max_iterations`` times) instead of removing skeleton points.""" + fitter = SkeletonSplineFitter() + fitter.ordered_coords = coords + fitter.fit_regularized_spline(coords, smoothing_factor=smoothing_factor, degree=degree) + fitter.sample_points_along_spline(spacing_distance) + smoothing = max(fitter.smoothing, float(len(coords))) + for _ in range(max_iterations): + if fitter.max_turn() <= curvature_threshold: + break + smoothing *= 2.0 + fitter.fit_regularized_spline(coords, smoothing_factor=smoothing, degree=degree) + fitter.sample_points_along_spline(spacing_distance) + transforms = fitter.compute_transforms() if compute_transforms else None + return fitter.sampled_points, transforms, fitter, fitter.get_spline_properties() def fit_spline_to_skeleton( @@ -328,91 +258,41 @@ def fit_spline_to_skeleton( max_iterations: int = 5, ) -> Tuple[np.ndarray, Optional[np.ndarray], SkeletonSplineFitter, Dict[str, Any]]: """ - Main function to fit a regularized 3D spline to a skeleton and sample points. + Fit a smoothing spline to the longest filament of a skeleton and sample points along it. Args: binary_volume: 3D binary volume where skeleton is True/1 - spacing_distance: Distance between consecutive sampled points along the spline - smoothing_factor: Smoothing parameter for spline fitting (auto if None) + spacing_distance: Distance between consecutive sampled points along the spline, in voxels + smoothing_factor: Smoothing parameter for spline fitting (scipy's ``s``, voxels squared; auto if None) degree: Degree of the spline (1-5) - connectivity_radius: Maximum distance to consider skeleton points as connected + connectivity_radius: Unused; skeleton voxels are connected to their 26 neighbours compute_transforms: Whether to compute 4x4 transformation matrices for each point - curvature_threshold: Maximum allowed curvature before outlier removal (default 0.2) - max_iterations: Maximum number of outlier removal iterations (default 5) + curvature_threshold: Largest sine of the turning angle between consecutive samples; a sharper fit is + smoothed further + max_iterations: Maximum number of smoothing increases Returns: Tuple of (sampled_points, transforms, spline_fitter, properties): - - sampled_points: Nx3 array of evenly spaced points along spline + - sampled_points: Nx3 array of points (x, y, z, voxels) exactly ``spacing_distance`` apart - transforms: [N, 4, 4] array of transformation matrices (or None if compute_transforms=False) - spline_fitter: SkeletonSplineFitter object for further analysis - properties: dict with spline properties """ - # Initialize fitter - fitter = SkeletonSplineFitter() - - # Extract skeleton coordinates - coords = fitter.extract_skeleton_coordinates(binary_volume) - - if len(coords) == 0: + chains = _skeleton_chains(binary_volume) + if not chains: raise ValueError("No skeleton points found in binary volume") - - # Order skeleton points - ordered_coords = fitter.order_skeleton_points_longest_path(coords, connectivity_radius=connectivity_radius) - - if len(ordered_coords) < 2: + coords = max((chain.coords for chain in chains), key=_chain_length) + if len(coords) < 2: raise ValueError("Not enough ordered points for spline fitting") - - # Iterative fitting with outlier removal - current_coords = ordered_coords.copy() - iteration = 0 - - while iteration < max_iterations: - # Fit regularized spline - fitter.fit_regularized_spline(current_coords, smoothing_factor=smoothing_factor, degree=degree) - - # Sample points along spline - sampled_points = fitter.sample_points_along_spline(spacing_distance) - - # Get properties to check curvature - properties = fitter.get_spline_properties() - max_curvature = properties.get("max_curvature", 0) - - print(f"Iteration {iteration + 1}: Max curvature = {max_curvature:.4f}") - - # If curvature is acceptable, break - if max_curvature <= curvature_threshold: - print(f"Curvature acceptable after {iteration + 1} iterations") - break - - # Remove outliers and try again - print(f"Max curvature {max_curvature:.4f} > {curvature_threshold}, removing outliers...") - filtered_coords = fitter.detect_high_curvature_outliers(current_coords, curvature_threshold) - - # If no points were removed, break to avoid infinite loop - if len(filtered_coords) == len(current_coords): - print("No outliers found to remove, stopping iterations") - break - - # If too few points remain, break - if len(filtered_coords) < degree + 1: - print(f"Too few points remaining ({len(filtered_coords)}), stopping iterations") - break - - current_coords = filtered_coords - iteration += 1 - - if iteration >= max_iterations: - print(f"Reached maximum iterations ({max_iterations}), final curvature: {max_curvature:.4f}") - - # Compute transformation matrices if requested - transforms = None - if compute_transforms: - transforms = fitter.compute_transforms() - - # Get final properties - properties = fitter.get_spline_properties() - - return sampled_points, transforms, fitter, properties + return _fit_chain( + coords, + spacing_distance, + smoothing_factor, + degree, + compute_transforms, + curvature_threshold, + max_iterations, + ) def fit_spline_to_segmentation( @@ -429,74 +309,108 @@ def fit_spline_to_segmentation( curvature_threshold: float = 0.2, max_iterations: int = 5, voxel_spacing: float = 1.0, + label: Optional[int] = None, + filaments_object_name: Optional[str] = None, + filaments_user_id: Optional[str] = None, + filaments_session_id: Optional[str] = None, + **kwargs, ) -> Optional[Tuple["CopickPicks", Dict[str, int]]]: """ - Fit a spline to a segmentation (skeleton) volume and create picks with orientations. + Fit a spline to every filament of a segmentation (skeleton) volume and create picks with orientations. Matches the lazy converter signature: (segmentation, run, object_name, session_id, user_id, **tool_kwargs) + Every filament gets its own spline. The picks of all filaments go into one pick set, grouped by filament and in + order along it, with the filament's ID (1..K by length, longest first) as ``instance_id``. + Args: segmentation: Input segmentation containing skeleton to fit spline to run: CopickRun object object_name: Name for the output pick object session_id: Session ID for output picks user_id: User ID for output picks - spacing_distance: Distance between consecutive sampled points along the spline - smoothing_factor: Smoothing parameter for spline fitting (auto if None) + spacing_distance: Distance between consecutive sampled points along the spline, in voxels + smoothing_factor: Smoothing parameter for spline fitting (scipy's ``s``, voxels squared; auto if None) degree: Degree of the spline (1-5) - connectivity_radius: Maximum distance to consider skeleton points as connected + connectivity_radius: Unused; skeleton voxels are connected to their 26 neighbours compute_transforms: Whether to compute orientations for picks - curvature_threshold: Maximum allowed curvature before outlier removal - max_iterations: Maximum number of outlier removal iterations - voxel_spacing: Voxel spacing for coordinate scaling + curvature_threshold: Largest sine of the turning angle between consecutive samples + max_iterations: Maximum number of smoothing increases + voxel_spacing: Voxel spacing for coordinate scaling (the segmentation's own voxel size is used when known) + label: Label to fit in a multilabel segmentation (default: every non-zero voxel) + filaments_object_name: Also store the fitted splines as copick Filaments under this object name + filaments_user_id: User ID of the Filaments + filaments_session_id: Session ID of the Filaments ({input_session_id} is replaced) + **kwargs: Additional keyword arguments from the lazy converter Returns: Tuple of (CopickPicks object, stats dict) or None if failed. - Stats dict contains 'picks_created'. + Stats dict contains 'picks_created' and 'filaments_fitted'. """ - # Get the segmentation volume volume = segmentation.numpy() if volume is None: logger.error(f"Could not load segmentation data for {run.name}") return None - logger.info(f"Fitting spline to segmentation {segmentation.session_id} in run {run.name}") + logger.info(f"Fitting splines to segmentation {segmentation.session_id} in run {run.name}") try: - # Fit spline to skeleton - sampled_points, transforms, fitter, properties = fit_spline_to_skeleton( - binary_volume=volume.astype(bool), - spacing_distance=spacing_distance, - smoothing_factor=smoothing_factor, - degree=degree, - connectivity_radius=connectivity_radius, - compute_transforms=compute_transforms, - curvature_threshold=curvature_threshold, - max_iterations=max_iterations, - ) - - # Scale points to physical coordinates - scaled_points = sampled_points * voxel_spacing - - logger.info(f"Spline properties: {properties}") - - # Create output picks - output_picks = run.new_picks( - object_name=object_name, - session_id=session_id, - user_id=user_id, - exist_ok=True, - ) - - # Store the picks with transformations - if compute_transforms and transforms is not None: - output_picks.from_numpy(scaled_points, transforms) + voxel_spacing = float(segmentation.voxel_size or voxel_spacing) + mask = volume == label if label is not None else volume > 0 + chains = [chain.coords for chain in _skeleton_chains(mask)] + chains = sorted((c for c in chains if len(c) >= 2), key=_chain_length, reverse=True) + + positions, transforms, ids, curves = [], [], [], [] + for instance_id, coords in enumerate(chains, start=1): + points, rotations, fitter, _ = _fit_chain( + _start_at_smaller_end(coords), + spacing_distance, + smoothing_factor, + degree, + True, + curvature_threshold, + max_iterations, + ) + positions.append(points * voxel_spacing) + transforms.append(rotations if compute_transforms else np.tile(np.eye(4), (len(points), 1, 1))) + ids.append(np.full(len(points), instance_id, dtype=np.int64)) + curves.append((fitter.spline_tck, fitter.smoothing)) + + output_picks = run.new_picks(object_name=object_name, session_id=session_id, user_id=user_id, exist_ok=True) + if positions: + output_picks.from_numpy( + np.concatenate(positions), + np.concatenate(transforms), + instance_ids=np.concatenate(ids), + ) else: - output_picks.from_numpy(scaled_points) - - stats = {"picks_created": len(scaled_points)} - logger.info(f"Created {stats['picks_created']} picks with session_id: {session_id}") + output_picks.points = [] + output_picks.store() + + if filaments_object_name: + _store_filaments( + run, + segmentation, + curves, + voxel_spacing, + filaments_object_name, + filaments_user_id or user_id, + (filaments_session_id or session_id).replace("{input_session_id}", segmentation.session_id), + { + "spacing_distance": spacing_distance, + "smoothing_factor": smoothing_factor, + "degree": degree, + "curvature_threshold": curvature_threshold, + "max_iterations": max_iterations, + }, + ) + + stats = {"picks_created": int(sum(len(p) for p in positions)), "filaments_fitted": len(positions)} + logger.info( + f"Created {stats['picks_created']} picks on {stats['filaments_fitted']} filaments " + f"with session_id: {session_id}", + ) return output_picks, stats except Exception as e: @@ -504,6 +418,40 @@ def fit_spline_to_segmentation( return None +def _start_at_smaller_end(coords: np.ndarray) -> np.ndarray: + """Order a chain so that it starts at the end with the smaller (z, y, x) voxel.""" + start, end = np.round(coords[0]).astype(int), np.round(coords[-1]).astype(int) + return coords[::-1] if tuple(end) < tuple(start) else coords + + +def _store_filaments(run, segmentation, curves, voxel_spacing, object_name, user_id, session_id, flags) -> None: + """Store fitted splines (voxel coordinates) as copick Filaments with ``bspline`` curves.""" + from copick.models import CopickFilamentCurve + + import copick_utils + + source = f"{segmentation.name}:{segmentation.user_id}/{segmentation.session_id}@{voxel_spacing}" + fitted = [ + CopickFilamentCurve.from_tck(tck, step=voxel_spacing, smoothing=smoothing, scale=voxel_spacing) + for tck, smoothing in curves + ] + metadata = [ + { + "fit": { + "method": "splprep", + "coordinates": "voxels", + "tool": "copick-utils fit-spline", + "version": copick_utils.__version__, + "source": source, + **flags, + }, + } + for _ in curves + ] + filaments = run.new_filaments(object_name, session_id, user_id, exist_ok=True) + filaments.from_curves(fitted, voxel_spacing=voxel_spacing, metadata=metadata) + + # Lazy batch converter for the lazy task discovery architecture fit_spline_lazy_batch = create_lazy_batch_converter( converter_func=fit_spline_to_segmentation, diff --git a/tests/conftest.py b/tests/conftest.py index 807e7c0..b3bc432 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,8 @@ "name": "microtubule", "is_particle": True, "label": 1, - "radius": 120.0, + # Tube radius: the synthetic test tubes have a 30 Å label radius, above a third of this (seg2fil's default). + "radius": 75.0, "metadata": {"copick": {"filament": {"polar": True}}}, }, {"name": "ribosome", "is_particle": True, "label": 2, "radius": 150.0}, diff --git a/tests/test_converter_signatures.py b/tests/test_converter_signatures.py index b94708b..d48e3ed 100644 --- a/tests/test_converter_signatures.py +++ b/tests/test_converter_signatures.py @@ -20,11 +20,14 @@ class of "signature drift" bugs (seg2mesh, mesh2picks, picks2{mesh,surface,ellip import pytest from copick_utils.converters.caps_from_mesh import caps_from_mesh_lazy_batch from copick_utils.converters.ellipsoid_from_picks import ellipsoid_from_picks_lazy_batch +from copick_utils.converters.filaments_from_segmentation import filaments_from_segmentation_lazy_batch from copick_utils.converters.mesh_from_picks import mesh_from_picks_lazy_batch from copick_utils.converters.mesh_from_segmentation import mesh_from_segmentation_lazy_batch +from copick_utils.converters.picks_from_filaments import picks_from_filaments_lazy_batch from copick_utils.converters.picks_from_mesh import picks_from_mesh_lazy_batch from copick_utils.converters.picks_from_segmentation import picks_from_segmentation_lazy_batch from copick_utils.converters.plane_from_picks import plane_from_picks_lazy_batch +from copick_utils.converters.segmentation_from_filaments import segmentation_from_filaments_lazy_batch from copick_utils.converters.segmentation_from_mesh import segmentation_from_mesh_lazy_batch from copick_utils.converters.segmentation_from_picks import segmentation_from_picks_lazy_batch from copick_utils.converters.sphere_from_picks import sphere_from_picks_lazy_batch @@ -52,6 +55,9 @@ class of "signature drift" bugs (seg2mesh, mesh2picks, picks2{mesh,surface,ellip ("picks2plane", plane_from_picks_lazy_batch, "picks", False), ("picks2sphere", sphere_from_picks_lazy_batch, "picks", False), ("mesh2caps", caps_from_mesh_lazy_batch, "mesh", False), + ("seg2fil", filaments_from_segmentation_lazy_batch, "segmentation", False), + ("fil2picks", picks_from_filaments_lazy_batch, "filaments", False), + ("fil2seg", segmentation_from_filaments_lazy_batch, "filaments", False), ("picksin", picks_inclusion_by_mesh_lazy_batch, "picks", True), ("picksout", picks_exclusion_by_mesh_lazy_batch, "picks", True), ("clippicks", limit_picks_by_distance_lazy_batch, "picks", True), diff --git a/tests/test_filaments.py b/tests/test_filaments.py new file mode 100644 index 0000000..896bc74 --- /dev/null +++ b/tests/test_filaments.py @@ -0,0 +1,382 @@ +"""Filament tracing: seg2fil, fil2picks, fil2seg, and the fit-spline and skeletonize fixes, on synthetic tubes.""" + +import numpy as np +import pytest +from click.testing import CliRunner +from copick.util.filaments import curve_is_current, evaluate_curve +from copick.util.relion import filament_track_lengths +from copick_utils.cli.fil2picks import fil2picks +from copick_utils.cli.fil2seg import fil2seg +from copick_utils.cli.fit_spline import fit_spline +from copick_utils.cli.seg2fil import seg2fil +from copick_utils.cli.skeletonize import skeletonize +from copick_utils.process.filament_tracing import ( + TraceParameters, + evaluate_curve_with_tangents, + evaluate_spline, + sample_curve, + trace_centrelines, +) + +VS = 10.0 +SHAPE = (40, 80, 80) + + +def tube(volume, start, end, radius, value=1): + """Paint a straight tube from ``start`` to ``end`` (z, y, x voxels).""" + start, end = np.asarray(start, float), np.asarray(end, float) + zz, yy, xx = np.indices(volume.shape) + grid = np.stack([zz, yy, xx], -1).astype(float) + d = end - start + t = np.clip(((grid - start) @ d) / (d @ d), 0, 1) + volume[np.linalg.norm(grid - (start + t[..., None] * d), axis=-1) <= radius] = value + return volume + + +def arc(volume, radius_px=25, tube_radius=3): + zz, yy, xx = np.indices(volume.shape) + r = np.hypot(yy - 40, xx - 40) + volume[(np.abs(r - radius_px) <= tube_radius) & (np.abs(zz - 20) <= tube_radius) & (xx >= 40)] = 1 + return volume + + +def _invoke(command, args): + result = CliRunner().invoke(command, args, catch_exceptions=False) + assert result.exit_code == 0, result.output + return result + + +def _write_seg(run, array, name="microtubule", session="s", **kwargs): + run.new_voxel_spacing(VS, exist_ok=True) + seg = run.new_segmentation(VS, name, session, user_id="test", exist_ok=True, **kwargs) + seg.from_numpy(array) + return seg + + +def _bspline(c): + knots, coefficients, k = c.tck + return {"kind": "bspline", "knots": knots, "control_points": np.stack(coefficients, 1), "degree": k, "step": VS} + + +# --- the engine -------------------------------------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "name, build, expected", + [ + ("straight", lambda v: tube(v, (20, 40, 10), (20, 40, 70), 3), 1), + ("along y", lambda v: tube(v, (20, 5, 40), (20, 75, 40), 3), 1), + ("arc", arc, 1), + ( + "crossing in one component", + lambda v: tube(tube(v, (20, 40, 5), (20, 40, 75), 3), (20, 5, 40), (20, 75, 40), 3), + 2, + ), + ("two tubes", lambda v: tube(tube(v, (20, 20, 5), (20, 20, 75), 3), (20, 60, 5), (20, 60, 75), 3), 2), + ], +) +def test_tracing_finds_each_filament(name, build, expected): + volume = build(np.zeros(SHAPE, np.uint8)) + centrelines, assignment, report = trace_centrelines(volume, VS, TraceParameters(min_length=150)) + assert len(centrelines) == expected + assert [c.instance_id for c in centrelines] == list(range(1, expected + 1)) + assert sorted(np.unique(assignment[volume > 0]).tolist()) == list(range(1, expected + 1)) + assert np.all(assignment[volume == 0] == 0) + for c in centrelines: + assert c.score > 0.95 # the fitted centreline stays inside the label + assert abs(c.radius - 30.0) < 6.0 + + +@pytest.mark.parametrize("seed", [0, 2, 6, 8]) +def test_a_thick_end_does_not_turn_the_filament_back(seed): + # A filament ending in a thick, rough blob (a splayed microtubule end on 10521) skeletonizes to a tangle of small + # loops that merge into one large junction; continuing "straight" through it must not leave by a branch that + # heads back alongside the filament. + volume = tube(np.zeros((40, 80, 100), np.uint8), (20, 40, 5), (20, 40, 70), 3.5) + zz, yy, xx = np.indices(volume.shape) + ellipsoid = ((zz - 20) / 6.5) ** 2 + ((yy - 41) / 9) ** 2 + ((xx - 76) / 13) ** 2 + volume[ellipsoid + np.random.default_rng(seed).normal(0, 0.15, volume.shape) <= 1] = 1 + centrelines, _, _ = trace_centrelines(volume, VS, TraceParameters(min_radius=20.0)) + assert len(centrelines) == 1 + points, tangents, _ = evaluate_spline(centrelines[0].tck, step=VS / 2) + overall = (points[-1] - points[0]) / np.linalg.norm(points[-1] - points[0]) + assert np.degrees(np.arccos(np.clip(tangents @ overall, -1, 1))).max() < 45.0 + + +def test_a_tube_labelled_by_its_wall_traces_as_one_filament(): + zz, yy, xx = np.indices(SHAPE) + r = np.hypot(zz - 20, yy - 40) + volume = ((r <= 6) & (r >= 3) & (xx >= 10) & (xx <= 70)).astype(np.uint8) # an empty lumen + centrelines, _, _ = trace_centrelines(volume, VS, TraceParameters(fill_lumen=60.0, min_radius=40.0)) + assert len(centrelines) == 1 and centrelines[0].length > 550 and centrelines[0].radius > 40.0 + unfilled, _, _ = trace_centrelines(volume, VS, TraceParameters(min_radius=40.0)) + assert len(unfilled) != 1 or unfilled[0].length < 550 + + +def test_specks_blobs_and_thin_slivers_are_rejected(): + volume = tube(np.zeros(SHAPE, np.uint8), (20, 40, 10), (20, 40, 70), 3) + volume[2:4, 2:4, 2:4] = 1 # speck + volume[28:38, 60:72, 60:72] = 1 # compact blob + tube(volume, (5, 10, 10), (5, 10, 60), 0.6) # thin sliver, label radius ~1 voxel + centrelines, _, report = trace_centrelines(volume, VS, TraceParameters(min_radius=20.0)) + assert len(centrelines) == 1 and centrelines[0].length > 500 + assert report.rejected_length + report.rejected_aspect >= 2 and report.rejected_radius == 1 + + +def test_y_junction_continues_the_straightest_pair(): + volume = tube(np.zeros(SHAPE, np.uint8), (20, 40, 5), (20, 40, 75), 3) # straight through + tube(volume, (20, 40, 40), (20, 70, 55), 3) # side branch at ~63 degrees + centrelines, _, report = trace_centrelines(volume, VS, TraceParameters(min_length=150)) + lengths = sorted(c.length for c in centrelines) + assert len(centrelines) == 2 and lengths[-1] > 650 # the through-filament is not split at the junction + assert report.junctions_resolved == 1 + + +def test_ring_is_one_closed_filament(): + volume = np.zeros(SHAPE, np.uint8) + zz, yy, xx = np.indices(SHAPE) + volume[(np.abs(np.hypot(yy - 40, xx - 40) - 25) <= 3) & (np.abs(zz - 20) <= 3)] = 1 + centrelines, _, _ = trace_centrelines(volume, VS) + assert len(centrelines) == 1 and centrelines[0].metadata["closed"] + assert abs(centrelines[0].length - 2 * np.pi * 250) < 60 + + +def test_sampling_spacing_frames_and_track_lengths(): + centrelines, _, _ = trace_centrelines(tube(np.zeros(SHAPE, np.uint8), (20, 5, 40), (20, 75, 40), 3), VS) + curve = _bspline(centrelines[0]) + positions, rotations, along = sample_curve(curve, 82.0) + np.testing.assert_allclose(np.diff(along), 82.0) + np.testing.assert_allclose(np.linalg.det(rotations), 1.0, atol=1e-9) + # +Z is the tangent in point order (the filament starts at its smaller-y end) + np.testing.assert_allclose(rotations[:, :, 2], np.tile([0.0, 1.0, 0.0], (len(rotations), 1)), atol=1e-3) + # The frames along Y (the tilt axis) never jump + x = rotations[:, :, 0] + assert np.max(np.degrees(np.arccos(np.clip(np.sum(x[1:] * x[:-1], axis=1), -1, 1)))) < 1.0 + np.testing.assert_allclose(filament_track_lengths(positions, np.ones(len(positions), int)), along - along[0]) + + +def test_bspline_tangent_matches_the_curve_direction(): + centrelines, _, _ = trace_centrelines(arc(np.zeros(SHAPE, np.uint8)), VS) + curve = _bspline(centrelines[0]) + points, tangents, _ = evaluate_curve_with_tangents(curve, 1.0) + chords = np.diff(points, axis=0) + chords /= np.linalg.norm(chords, axis=1, keepdims=True) + assert np.min(np.sum(chords * tangents[:-1], axis=1)) > 0.999 + + +# --- the commands ------------------------------------------------------------------------------------------------------ + + +def _traced(run, config_path, volume): + _write_seg(run, volume) + _invoke( + seg2fil, + [ + "-c", + config_path, + "-i", + "microtubule:test/s@10.0", + "-o", + "microtubule:trace/s", + "--instances", + "microtubule:trace/s@10.0?instance=true", + "--min-length", + "150", + ], + ) + run.refresh() + filaments = run.get_filaments("microtubule", "trace", "s")[0] + instances = run.get_segmentations(name="microtubule", user_id="trace", session_id="s", is_instance=True)[0] + return filaments, instances + + +def test_ids_are_shared_by_instances_filaments_and_picks(run, config_path): + volume = tube(tube(np.zeros(SHAPE, np.uint8), (20, 40, 5), (20, 40, 75), 3), (20, 5, 40), (20, 75, 40), 3) + filaments, instances = _traced(run, config_path, volume) + ids = sorted(f.instance_id for f in filaments.filaments) + assert ids == [1, 2] + assert sorted(np.unique(instances.numpy()[volume > 0]).tolist()) == ids + + # Each stored curve is the exact fit, and copick's points are its evaluation + for f in filaments.filaments: + assert f.curve.kind == "bspline" and f.curve_is_current() + regenerated = evaluate_curve( + f.curve.control_points, + f.curve.step, + "bspline", + None, + f.curve.degree, + f.curve.knots, + ) + np.testing.assert_allclose(np.asarray(f.points), regenerated, atol=1e-6) + assert f.metadata["fit"]["method"] == "splprep" and f.metadata["fit"]["tool"] == "copick-utils seg2fil" + + _invoke(fil2picks, ["-c", config_path, "-i", "microtubule:trace/s", "-o", "microtubule:trace/s", "--spacing", "82"]) + run.refresh() + picks = run.get_picks("microtubule", "trace", "s")[0] + pick_ids = picks.instance_ids() + assert sorted(set(pick_ids.tolist())) == ids + assert np.all(np.diff(pick_ids) >= 0) # grouped by filament + # every pick sits inside its own filament's voxels + zyx = np.round(picks.full_positions()[:, ::-1] / VS).astype(int) + assert np.all(instances.numpy()[tuple(zyx.T)] == pick_ids) + + +def test_retracing_the_instance_segmentation_reproduces_and_honours_edits(run, config_path): + volume = tube(tube(np.zeros(SHAPE, np.uint8), (20, 20, 5), (20, 20, 75), 3), (20, 60, 5), (20, 60, 75), 3) + filaments, instances = _traced(run, config_path, volume) + first = {f.instance_id: np.asarray(f.points) for f in filaments.filaments} + + _invoke( + seg2fil, + ["-c", config_path, "-i", "microtubule:trace/s@10.0?instance=true", "-o", "microtubule:retrace/s"], + ) + run.refresh() + again = {f.instance_id: np.asarray(f.points) for f in run.get_filaments("microtubule", "retrace", "s")[0].filaments} + assert set(again) == set(first) + for fid in first: + assert np.max(np.min(np.linalg.norm(again[fid][:, None] - first[fid][None], axis=-1), axis=1)) < VS + + # An edit (renumber filament 2 to 7, delete filament 1) is what the re-trace returns + edited = instances.numpy().astype(np.uint16) + edited[edited == 1] = 0 + edited[edited == 2] = 7 + _write_seg(run, edited, session="edited", is_instance=True) + _invoke( + seg2fil, + ["-c", config_path, "-i", "microtubule:test/edited@10.0?instance=true", "-o", "microtubule:retrace/edited"], + ) + run.refresh() + assert [f.instance_id for f in run.get_filaments("microtubule", "retrace", "edited")[0].filaments] == [7] + + +def test_hand_traced_and_plain_filaments_sample_smoothly(run, config_path): + run.new_voxel_spacing(VS, exist_ok=True) + clicks = np.array([[100.0, 100.0, 200.0], [300.0, 250.0, 200.0], [500.0, 300.0, 210.0], [700.0, 260.0, 220.0]]) + hand = run.new_filaments("microtubule", "hand", "test", exist_ok=True) + hand.from_control_points([clicks], voxel_spacing=VS) + plain = run.new_filaments("microtubule", "plain", "test", exist_ok=True) + plain.from_numpy([np.asarray(hand.filaments[0].points)], voxel_spacing=VS) + + for session in ("hand", "plain"): + _invoke( + fil2picks, + [ + "-c", + config_path, + "-i", + f"microtubule:test/{session}", + "-o", + f"microtubule:picks/{session}", + "--spacing", + "40", + ], + ) + run.refresh() + picks = run.get_picks("microtubule", "picks", session)[0] + _, transforms = picks.numpy() + z = transforms[:, :3, 2] + chords = np.diff(picks.full_positions(), axis=0) + chords /= np.linalg.norm(chords, axis=1, keepdims=True) + assert np.min(np.sum(z[:-1] * chords, axis=1)) > 0.95 # +Z follows the filament + assert np.min(np.sum(z[1:] * z[:-1], axis=1)) > 0.95 # and turns smoothly + + +def test_reversed_filament_gives_reversed_picks(): + centrelines, _, _ = trace_centrelines(tube(np.zeros(SHAPE, np.uint8), (20, 40, 10), (20, 40, 70), 3), VS) + forward = _bspline(centrelines[0]) + knots = np.asarray(forward["knots"]) + backward = dict(forward, knots=(knots[0] + knots[-1] - knots[::-1]), control_points=forward["control_points"][::-1]) + p1, r1, _ = sample_curve(forward, 82.0) + p2, r2, _ = sample_curve(backward, 82.0) + np.testing.assert_allclose(p1, p2[::-1], atol=1e-6) + np.testing.assert_allclose(r1[:, :, 2], -r2[::-1, :, 2], atol=1e-6) + + +def test_fil2seg_paints_tubes_that_retrace_to_the_filaments(run, config_path): + volume = tube(np.zeros(SHAPE, np.uint8), (20, 40, 10), (20, 40, 70), 3) + filaments, _ = _traced(run, config_path, volume) + run.get_voxel_spacing(VS).new_tomogram("wbp").from_numpy(np.zeros(SHAPE, np.float32)) + _invoke(fil2seg, ["-c", config_path, "-i", "microtubule:trace/s", "-o", "microtubule:tubes/s@10.0"]) + run.refresh() + painted = run.get_segmentations(name="microtubule", user_id="tubes", session_id="s", is_instance=True)[0].numpy() + assert sorted(np.unique(painted).tolist()) == [0, 1] + centrelines, _, _ = trace_centrelines(painted, VS, instances=True) + original = np.asarray(filaments.filaments[0].points) + # Every original point lies on the re-traced centreline (which runs on into the tube's rounded caps) + assert np.max(np.min(np.linalg.norm(original[:, None] - centrelines[0].points[None], axis=-1), axis=1)) < VS + + +# --- the kept commands ------------------------------------------------------------------------------------------------- + + +def test_fit_spline_fits_every_filament_and_can_store_them(run, config_path): + volume = tube(tube(np.zeros(SHAPE, np.uint8), (20, 20, 5), (20, 20, 75), 3), (20, 60, 5), (20, 60, 75), 3) + _write_seg(run, volume) + _invoke( + fit_spline, + [ + "-c", + config_path, + "-i", + "microtubule:test/s@10.0", + "-o", + "microtubule:spline/s", + "--spacing-distance", + "8", + "--filaments", + "microtubule:spline/s", + ], + ) + run.refresh() + picks = run.get_picks("microtubule", "spline", "s")[0] + ids = picks.instance_ids() + assert sorted(set(ids.tolist())) == [1, 2] + for fid in (1, 2): + np.testing.assert_allclose(np.linalg.norm(np.diff(picks.full_positions()[ids == fid], axis=0), axis=1), 80.0) + filaments = run.get_filaments("microtubule", "spline", "s")[0] + assert len(filaments.filaments) == 2 + for f in filaments.filaments: + assert f.curve.kind == "bspline" and f.metadata["fit"]["coordinates"] == "voxels" + assert curve_is_current( + f.points, + f.curve.control_points, + f.curve.step, + "bspline", + f.curve.degree, + f.curve.knots, + ) + # the picks lie on the stored curve + points = np.asarray(f.points) + mine = picks.full_positions()[ids == f.instance_id] + assert np.max(np.min(np.linalg.norm(mine[:, None] - points[None], axis=-1), axis=1)) < 1.0 + + +def test_skeletonize_prunes_spurs_and_keeps_the_output_name(run, config_path): + volume = np.zeros(SHAPE, np.uint8) + volume[20, 40, 5:75] = 1 + volume[20, 41:44, 30] = 1 # a 3-voxel spur + _write_seg(run, volume) + _invoke( + skeletonize, + [ + "-c", + config_path, + "-i", + "microtubule:test/s@10.0", + "-o", + "membrane:skel/s@10.0", + "--prune-length", + "50", + "--keep-noise", + ], + ) + run.refresh() + out = run.get_segmentations(name="membrane", user_id="skel", session_id="s") + assert len(out) == 1 + skeleton = out[0].numpy() > 0 + # Thinning moves the junction one voxel into the spur; pruning removes the rest of the spur + assert skeleton[20, 42:44, 30].sum() == 0 + from scipy.ndimage import label + + assert label(skeleton, structure=np.ones((3, 3, 3)))[1] == 1 and skeleton.sum() >= 68 diff --git a/tests/test_flag_normalization.py b/tests/test_flag_normalization.py index 92c09e8..c6aebac 100644 --- a/tests/test_flag_normalization.py +++ b/tests/test_flag_normalization.py @@ -55,3 +55,14 @@ def test_fit_spline_derives_vs_from_uri(runner): param, hidden = _hidden_alias(fit_spline, "--voxel-spacing") assert param is not None and hidden is True + + +def test_fit_spline_hides_the_retired_connectivity_radius(runner): + from copick_utils.cli.fit_spline import fit_spline + + out = runner.invoke(fit_spline, ["--help"]).output + assert "--connectivity-radius" not in out + assert "--filaments" in out and "--label" in out + + param, hidden = _hidden_alias(fit_spline, "--connectivity-radius") + assert param is not None and hidden is True diff --git a/tests/test_run_names_option.py b/tests/test_run_names_option.py index 4e2d2a8..a5bfd5d 100644 --- a/tests/test_run_names_option.py +++ b/tests/test_run_names_option.py @@ -24,6 +24,9 @@ def runner(): ("copick_utils.cli.segop", "segop"), ("copick_utils.cli.rescale", "rescale"), ("copick_utils.cli.combine_labels", "combine"), + ("copick_utils.cli.seg2fil", "seg2fil"), + ("copick_utils.cli.fil2picks", "fil2picks"), + ("copick_utils.cli.fil2seg", "fil2seg"), ]