Add explicit Z slice skipping

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
     )

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
@@ -260,13 +269,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 +307,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,
@@ -322,8 +343,11 @@ def _compute_flow(dataset,
 
         # 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, 
@@ -354,7 +378,7 @@ def _compute_flow(dataset,
             # 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, 
@@ -384,6 +408,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 +448,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

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)