From 76624ca49623b3053abd8b3df08220f05769d038 Mon Sep 17 00:00:00 2001 From: Marcel Sayre <52462209+msayr@users.noreply.github.com> Date: Sat, 4 Jul 2026 18:12:15 +1000 Subject: [PATCH 1/4] Add explicit Z slice skipping --- emalign/align_dataset_z.py | 23 ++++++++++++++++---- emalign/scripts/align_stack_z.py | 37 ++++++++++++++++++++++++++++++-- 2 files changed, 54 insertions(+), 6 deletions(-) diff --git a/emalign/align_dataset_z.py b/emalign/align_dataset_z.py index 03f09cf..d82967e 100644 --- a/emalign/align_dataset_z.py +++ b/emalign/align_dataset_z.py @@ -203,7 +203,7 @@ def initialize_destination_stores(destination_path, align_plan, save_downsampled return destination, destination_mask, ds_destination, ds_project_output_path -def execute_alignment(paths, dataset_configs, root_stack, num_workers, wipe_progress_stacks): +def execute_alignment(paths, dataset_configs, root_stack, num_workers, wipe_progress_stacks, skip_slices_global=None): '''Execute the alignment for all datasets. Args: @@ -212,7 +212,9 @@ def execute_alignment(paths, dataset_configs, root_stack, num_workers, wipe_prog root_stack: Name of the root stack num_workers: Number of worker threads wipe_progress_stacks: Optional stack name to wipe progress for + skip_slices_global: Optional project/global Z slices to skip for every stack ''' + skip_slices_global = sorted(set(skip_slices_global or [])) for i, path in enumerate(paths): for dataset_name in path: if dataset_name not in dataset_configs: @@ -226,6 +228,10 @@ def execute_alignment(paths, dataset_configs, root_stack, num_workers, wipe_prog config['num_workers'] = num_workers config['wipe_progress_flag'] = any([dataset_name == s for s in wipe_progress_stacks]) + if skip_slices_global: + config['skip_slices_global'] = sorted( + set(config.get('skip_slices_global', [])) | set(skip_slices_global) + ) # Start alignment try: @@ -247,7 +253,8 @@ def align_dataset_z(project_dir: str, num_workers: int = NUM_WORKERS, save_downsampled: float = DOWNSAMPLE_SCALE, start_over: bool = False, - wipe_progress_stacks: Optional[list[str]] = None) -> None: + wipe_progress_stacks: Optional[list[str]] = None, + skip_slices_global: Optional[list[int]] = None) -> None: '''Execute Z alignment using pre-generated configuration files. Args: @@ -256,6 +263,7 @@ def align_dataset_z(project_dir: str, save_downsampled: Downsampling factor for inspection store start_over: Wipe all progress and restart wipe_progress_stacks: List of specific stacks to wipe progress for + skip_slices_global: Project/global Z slice indices to skip ''' config_dir = os.path.join(project_dir, 'config/z_config') if not os.path.exists(config_dir) or not os.listdir(config_dir): @@ -311,7 +319,7 @@ def align_dataset_z(project_dir: str, # Execute alignment logging.info('Starting Z alignment...') logging.info(f'Number of cores used for rendering: {num_workers}') - execute_alignment(paths, dataset_configs, root_stack, num_workers, wipe_progress_stacks) + execute_alignment(paths, dataset_configs, root_stack, num_workers, wipe_progress_stacks, skip_slices_global) logging.info('Done!') logging.info(f'Output: {destination_path}') @@ -358,6 +366,12 @@ def align_dataset_z(project_dir: str, nargs='+', default=[''], help='Wipe progress for one or more specific stack(s) before starting') + parser.add_argument('--skip-slices', + dest='skip_slices_global', + type=int, + nargs='+', + default=[], + help='Project/global Z slice indices to skip during Z alignment and rendering') args = parser.parse_args() @@ -375,5 +389,6 @@ def align_dataset_z(project_dir: str, num_workers=args.num_workers, save_downsampled=args.save_downsampled, start_over=args.start_over, - wipe_progress_stacks=args.wipe_progress_stacks + wipe_progress_stacks=args.wipe_progress_stacks, + skip_slices_global=args.skip_slices_global ) diff --git a/emalign/scripts/align_stack_z.py b/emalign/scripts/align_stack_z.py index 75b9347..ffb407b 100644 --- a/emalign/scripts/align_stack_z.py +++ b/emalign/scripts/align_stack_z.py @@ -50,6 +50,8 @@ def align_stack_z(destination_path, local_z_max=None, xy_offset=[0,0], ignore_slices_flow=[], + skip_slices_global=[], + skip_slices_local=[], save_downsampled=1, overwrite=False, wipe_progress_flag=False, @@ -123,6 +125,24 @@ def align_stack_z(destination_path, dataset_mask = open_store(os.path.abspath(dataset_path) + '_mask', mode='r', dtype=ts.bool, allow_missing=True) dataset_mask = dataset_mask[dataset.domain] if dataset_mask is not None else dataset_mask + # Slices can be skipped either in the dataset's local Z coordinates or in + # final/global project coordinates. Flow computation and rendering both + # operate on the local TensorStore indices, so normalize everything here. + skip_slices_local = set(skip_slices_local or []) + skip_slices_global = set(skip_slices_global or []) + ignore_slices_flow = set(ignore_slices_flow or []) + normalized_skip_slices = ignore_slices_flow | skip_slices_local + normalized_skip_slices |= { + z - z_offset + dataset.domain.inclusive_min[0] + for z in skip_slices_global + } + normalized_skip_slices = sorted( + z for z in normalized_skip_slices + if dataset.domain.inclusive_min[0] <= z < dataset.domain.exclusive_max[0] + ) + if normalized_skip_slices: + logging.info(f'{dataset_name}: User-requested skipped local slices: {normalized_skip_slices}') + # Make the resolution match the target res = attrs['resolution'][-1] if yx_target_resolution is None: @@ -205,7 +225,7 @@ def align_stack_z(destination_path, # Compute flow and save to file flow, transform, bbox_ref = compute_flow_dataset(dataset=dataset, original_shape=original_shape, - ignore_slices=ignore_slices_flow, + ignore_slices=normalized_skip_slices, scale=scale, patch_size=patch_size, stride=stride, @@ -369,9 +389,16 @@ def align_stack_z(destination_path, skipped += 1 # Assume skipped slices are logged correctly continue + global_z = z + z_offset - dataset.domain.inclusive_min[0] + + if z in normalized_skip_slices: + skipped += 1 + metadata = {'skipped': True, 'user_skipped': True, 'empty_slice': False} + log_progress(db, dataset_name, step_name, global_z, z, metadata) + continue + # Load data data = dataset[z].read().result() - global_z = z + z_offset - dataset.domain.inclusive_min[0] if not data.any(): # If empty slice, skip and go to next z @@ -486,6 +513,10 @@ def align_stack_z(destination_path, parser = argparse.ArgumentParser(description='Align a stack in Z.') parser.add_argument('config_file', type=str, help='Path to the JSON configuration file.') parser.add_argument('--wipe-progress', action='store_true', help='Wipe progress for the specified stack before starting.') + parser.add_argument('--skip-slices-global', nargs='+', type=int, default=[], + help='Project/global Z slice indices to skip during Z alignment and rendering.') + parser.add_argument('--skip-slices-local', nargs='+', type=int, default=[], + help='Local input-stack Z slice indices to skip during Z alignment and rendering.') args = parser.parse_args() @@ -501,5 +532,7 @@ def align_stack_z(destination_path, config['project_name'] = project_name config['mongodb_config_filepath'] = mongodb_config_filepath config['wipe_progress_flag'] = args.wipe_progress + config['skip_slices_global'] = sorted(set(config.get('skip_slices_global', [])) | set(args.skip_slices_global)) + config['skip_slices_local'] = sorted(set(config.get('skip_slices_local', [])) | set(args.skip_slices_local)) align_stack_z(**config) From 38c2052dfd17b6949398a4afc19bf0f1a1c5d7ad Mon Sep 17 00:00:00 2001 From: Marcel Sayre <52462209+msayr@users.noreply.github.com> Date: Sun, 5 Jul 2026 01:20:51 +1000 Subject: [PATCH 2/4] Handle leading skipped Z slices --- emalign/align_z/align_z.py | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/emalign/align_z/align_z.py b/emalign/align_z/align_z.py index 61c3d3d..c435b9d 100644 --- a/emalign/align_z/align_z.py +++ b/emalign/align_z/align_z.py @@ -260,13 +260,25 @@ def _compute_flow(dataset, #---------- Start processing ----------# mfc = flow_field.JAXMaskedXCorrWithStatsCalculator() + pending_invalid_flows = 0 + + def append_invalid_flow(): + nonlocal pending_invalid_flows + if flows: + flows.append(np.ones_like(flows[-1]) * np.nan) + else: + # Flow shape is only known after the first valid pair is computed. + # Keep leading skipped/empty slices pending and backfill them once + # a real flow exists instead of indexing flows[-1]. + pending_invalid_flows += 1 + pbar = tqdm(range(start, dataset.domain.exclusive_max[0]), position=0, dynamic_ncols=True) for z in pbar: if z in ignore_slices: pbar.set_description(f'{dataset_name}: Ignoring slice...') # Slice is to be ignored for flow computation based on user input. # These should not be used for mesh relaxation or they will bias the result, so we set them as invalid. - flows.append(np.ones_like(flows[-1]) * np.nan) + append_invalid_flow() metadata = { 'ref_dataset': ref_dataset_name, @@ -286,7 +298,7 @@ def _compute_flow(dataset, # If empty slice, skip and compare to next one if not mov.any(): # We should be starting with a non-empty slice, so by the time we hit this, flow should exist - flows.append(np.ones_like(flows[-1]) * np.nan) + append_invalid_flow() metadata = { 'ref_dataset': ref_dataset_name, 'scale': scale, @@ -384,6 +396,9 @@ def _compute_flow(dataset, # Compute flow flow = _compute_flow_slice(overlap_ref, overlap_ref_mask, mov, mov_mask, mfc, patch_size, stride) + if pending_invalid_flows: + flows.extend([np.ones_like(flow) * np.nan] * pending_invalid_flows) + pending_invalid_flows = 0 flows.append(flow) # Save to file + database @@ -421,6 +436,11 @@ def _compute_flow(dataset, jax.clear_caches() + if pending_invalid_flows: + raise ValueError( + f'{dataset_name}: no valid slices were available after applying ignored/empty slices.' + ) + flows = homogenize_arrays_shape(flows, pad_value=np.nan) flows = np.transpose(flows, [1, 0, 2, 3]) # [channels, z, y, x] return flows, transform, bbox_ref, bbox_anchor From b66d059bd326e63b082f5c0b37339237c49b5b77 Mon Sep 17 00:00:00 2001 From: Marcel Sayre <52462209+msayr@users.noreply.github.com> Date: Sun, 5 Jul 2026 01:59:40 +1000 Subject: [PATCH 3/4] Set Z anchor after leading skips --- emalign/align_z/align_z.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/emalign/align_z/align_z.py b/emalign/align_z/align_z.py index c435b9d..3044695 100644 --- a/emalign/align_z/align_z.py +++ b/emalign/align_z/align_z.py @@ -334,8 +334,11 @@ def append_invalid_flow(): # Transform mov to match ref if transformations is None: - if z == anchor_z: - # This is the first slice, use the anchor bbox + if z == anchor_z or bbox_anchor is None: + # This is the first slice, or the first valid slice after one + # or more leading ignored/empty slices. Use/update the anchor + # bbox here so later downsampled flow computation has a valid + # bbox_anchor to scale. overlap_ref, overlap_ref_mask, bbox_anchor = get_overlap_ref(ref, mov, ref_mask=ref_mask, @@ -366,7 +369,7 @@ def append_invalid_flow(): # This gets added at the end of the array output_shape = np.array(output_shape) + patch_size else: - if z == anchor_z: + if z == anchor_z or bbox_anchor is None: overlap_ref, overlap_ref_mask, bbox_anchor = get_overlap_ref(ref, mov, ref_mask=ref_mask, From db44a0e08b65dccef67d0c55ae04599468caf6ab Mon Sep 17 00:00:00 2001 From: Marcel Sayre <52462209+msayr@users.noreply.github.com> Date: Sun, 5 Jul 2026 03:08:03 +1000 Subject: [PATCH 4/4] Resume Z flow with skipped anchor progress --- emalign/align_z/align_z.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/emalign/align_z/align_z.py b/emalign/align_z/align_z.py index 3044695..53cabc6 100644 --- a/emalign/align_z/align_z.py +++ b/emalign/align_z/align_z.py @@ -199,18 +199,27 @@ def _compute_flow(dataset, if check_progress(db, dataset_name, step_name, z, doc_filter={'scale': scale}): flows.append(dataset_flow[z].read().result()) transform[z] = dataset_trsf[z].read().result() if transformations is None else transformations[z] + progress_doc = db[dataset_name].find_one( + {'step_name': step_name, 'local_slice': z, 'scale': scale}, + {'bbox_anchor': 1, 'bbox_ref': 1} + ) or {} if z == anchor_z and bbox_anchor is None: # Get bbox from previous slices (should be passed but take it for return) - bbox_anchor = db[dataset_name].find_one({'step_name': step_name, 'local_slice': z, 'scale': scale}, - {'bbox_anchor': 1})['bbox_anchor'] + bbox_anchor = progress_doc.get('bbox_anchor') or progress_doc.get('bbox_ref') elif z > anchor_z and bbox_ref is None: # Get bbox from previous slices - bbox_ref = db[dataset_name].find_one({'step_name': step_name, 'local_slice': z, 'scale': scale}, - {'bbox_ref': 1})['bbox_ref'] + bbox_ref = progress_doc.get('bbox_ref') + if z > anchor_z and bbox_anchor is None: + # Older progress docs for leading skipped slices may not have + # bbox_anchor. Fall back to the first available bbox from a + # processed non-skipped slice so resume can continue. + bbox_anchor = progress_doc.get('bbox_anchor') or progress_doc.get('bbox_ref') if len(flows) == (dataset.domain.exclusive_max[0] - start): # Everything appears to have been processed, early exit + if bbox_anchor is None: + bbox_anchor = bbox_ref flows = homogenize_arrays_shape(flows, pad_value=np.nan) flows = np.transpose(flows, [1, 0, 2, 3]) # [channels, z, y, x] return flows, transform, bbox_ref, bbox_anchor