Recompute fuse XY configs when overwriting

emalign/align_xy/prep.py

@@ -223,7 +223,9 @@ def create_configs_fused_stacks(main_config_path,
         target_res = json.load(f)['resolution'][-1]
 
     # Find datasets
-    datasets, z_offsets = get_ordered_datasets([main_config_path], exclude=['flow', 'mask', '10x'])
+    datasets, z_offsets = get_ordered_datasets([main_config_path], exclude=['flow', 'mask', '10x', 'fused'])
+    dataset_names = [os.path.basename(os.path.abspath(ds.kvstore.path)) for ds in datasets]
+    logging.info('Found %d source dataset(s) for XY fusion: %s', len(dataset_names), dataset_names)
     z_ranges = [np.arange(z[0], z[0] + ds.shape[0]) for z, ds in zip(z_offsets, datasets)]
 
     # Find all ranges over which there is overlap
@@ -241,6 +243,13 @@ def create_configs_fused_stacks(main_config_path,
     for _, group in df.groupby('group'):
         z = group.z.min()
         indices = np.unique(group.ds_indices.to_numpy())[0]
+        group_names = [dataset_names[i] for i in indices]
+        logging.info(
+            'Testing XY fusion candidates for z=%s-%s: %s',
+            z,
+            group.z.max() + 1,
+            group_names,
+        )
 
         images = []
         for i in indices:
@@ -258,8 +267,21 @@ def create_configs_fused_stacks(main_config_path,
             valid_estimate = estimate_transform_sift(images[i], images[j], scale, refine_estimate=True)[3]
             if valid_estimate:
                 G.add_edge(indices[i], indices[j])
+                logging.info('Valid XY overlap: %s <-> %s', dataset_names[indices[i]], dataset_names[indices[j]])
+            else:
+                logging.warning('No valid XY overlap: %s <-> %s', dataset_names[indices[i]], dataset_names[indices[j]])
 
         # Valid matches are chained in case there are more than 2 matches for a range
+        matched_indices = set(G.nodes)
+        unmatched_names = [dataset_names[i] for i in indices if i not in matched_indices]
+        if unmatched_names:
+            logging.warning(
+                'No valid XY overlap was found for %s in z=%s-%s; these dataset(s) will not be included '
+                'in a fuse_xy config. This decision is made from SIFT matches on the image data, not from MongoDB.',
+                unmatched_names,
+                z,
+                group.z.max() + 1,
+            )
         for cc in nx.connected_components(G):
             config = {
                 'dataset_paths': [datasets[i].kvstore.path for i in cc], 
@@ -268,4 +290,4 @@ def create_configs_fused_stacks(main_config_path,
                 'zmax': int(group.z.max()) + 1 # Exclusive max
                 } 
             fused_configs.append(config)
-    return fused_configs
+    return fused_configs

emalign/align_z/utils.py

@@ -58,6 +58,9 @@ def get_ordered_datasets(config_paths, exclude=[]):
                 if any(check) or os.path.abspath(ds).endswith('_mask'):
                     # Always exclude masks from query
                     continue
+                if not os.path.exists(os.path.join(ds, '.zattrs')):
+                    logging.warning('Skipping incomplete dataset without .zattrs: %s', ds)
+                    continue
                 dataset = open_store(ds, mode='r')
                 z_shapes.append(dataset.shape[0])
 
@@ -372,4 +375,4 @@ def determine_initial_offset_ref(datasets, z_offsets, reference_path, reference_
         global_offset = np.min([global_offset, np.abs(xy_offset[::-1])], axis=0)  
         ref_bboxes.append(bbox_ref)
 
-    return np.abs(global_offset), ref_bboxes
+    return np.abs(global_offset), ref_bboxes

emalign/scripts/fuse_stacks_xy.py

@@ -2,8 +2,16 @@
 import json
 import logging
 import os
+import sys
 import traceback
 from datetime import datetime, UTC
+
+# Allow this file to be executed directly from a source checkout, e.g.
+# `python emalign/scripts/fuse_stacks_xy.py`, before importing the `emalign`
+# package.
+if __package__ in (None, ''):
+    sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
+
 from emalign.align_xy.stitch_offgrid import stitch_images
 from emalign.io.progress import get_mongo_client, get_mongo_db, log_progress, wipe_progress
 from emalign.io.store import write_data, open_store
@@ -78,13 +86,15 @@ def has_completed_fuse_progress(db, collection_name, step_name, local_slice_inde
 
 def get_fused_configs(
         main_config_path,
-        scale=0.1
+        scale=0.1,
+        recompute=False
         ):
     '''Gather or compute configuration files for groups of stacks to fuse.
 
     Args:
         config_path (str): Absolute path to the main_config.json file for this project.
         scale (float, optional): Scale to downsample images for determining offset using SIFT. Defaults to 0.1.
+        recompute (bool, optional): Whether to remove existing fuse_xy*.json files and recompute them.
 
     Returns:
         fused_configs (list of `dict`): list of configuration file per segment of stacks to fuse.
@@ -93,9 +103,17 @@ def get_fused_configs(
     # Output directory for the config files
     output_dir = os.path.dirname(os.path.abspath(main_config_path))
 
-    # Check for existing files
+    # Check for existing files.  These files are derived from the aligned
+    # xy_intermediate datasets, so they can become stale if stacks are added,
+    # removed, or re-run after the first fuse_stacks_xy invocation.
     config_filepaths = glob(os.path.join(output_dir, 'fuse_xy*.json'))
 
+    if recompute and config_filepaths:
+        logging.info('Removing %d existing fuse_xy config file(s) before recomputing.', len(config_filepaths))
+        for filepath in config_filepaths:
+            os.remove(filepath)
+        config_filepaths = []
+
     if len(config_filepaths) == 0:
         # Compute and write configuration files
         overlapping_groups = create_configs_fused_stacks(main_config_path, scale=scale)
@@ -108,6 +126,10 @@ def get_fused_configs(
                     json.dump(config, f, indent='')
     else:
         # Load configuration files
+        logging.warning(
+            'Loading existing fuse_xy config file(s). If xy_intermediate stacks changed since '
+            'these were generated, rerun with --recompute-configs to recompute them.'
+        )
         overlapping_groups = []
         pbar = tqdm(config_filepaths, position=0, desc='Loading existing configurations')
         for filepath in pbar:
@@ -362,6 +384,7 @@ def align_fused_stacks_xy(config_path,
                           stride=40,
                           img_on_top='auto',
                           overwrite=False,
+                          recompute_configs=False,
                           wipe_progress_stack=None,
                           num_workers=1):
     '''Align groups of overlapping stacks one after the other.
@@ -373,6 +396,7 @@ def align_fused_stacks_xy(config_path,
         stride (int, optional): _description_. Defaults to 40.
         img_on_top (str, optional): _description_. Defaults to 'auto'.
         overwrite (bool, optional): _description_. Defaults to False.
+        recompute_configs (bool, optional): Whether to remove and rebuild cached fuse_xy*.json files. Defaults to False.
         wipe_progress_stack (str, optional): Name of the stack to wipe progress for. Defaults to None.
         num_workers (int, optional): _description_. Defaults to 1.
     '''
@@ -388,7 +412,8 @@ def align_fused_stacks_xy(config_path,
 
 
     fused_configs = get_fused_configs(config_path,
-                                      0.1)
+                                      0.1,
+                                      recompute=recompute_configs)
 
     # Function to determine image quality to choose which one is on top
     # Highest value == on top
@@ -436,6 +461,9 @@ def align_fused_stacks_xy(config_path,
                         type=int,
                         help='Number of threads to use for rendering. Default: 1')
     parser.add_argument('--overwrite', action='store_true', help='Overwrite existing dataset.')
+    parser.add_argument('--recompute-configs',
+                        action='store_true',
+                        help='Delete existing fuse_xy*.json files and recompute them from current xy_intermediate datasets.')
     parser.add_argument('--wipe-progress',
                         dest='wipe_progress_stack',
                         type=str,
@@ -447,4 +475,5 @@ def align_fused_stacks_xy(config_path,
     align_fused_stacks_xy(config_path=args.config_path,
                           num_workers=args.num_workers,
                           overwrite=args.overwrite,
-                          wipe_progress_stack=args.wipe_progress_stack)
+                          recompute_configs=args.recompute_configs,
+                          wipe_progress_stack=args.wipe_progress_stack)