diff --git a/emalign/align_xy/prep.py b/emalign/align_xy/prep.py index 3435650..e51bff6 100644 --- a/emalign/align_xy/prep.py +++ b/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 \ No newline at end of file + return fused_configs diff --git a/emalign/align_z/utils.py b/emalign/align_z/utils.py index 18c523c..3945911 100644 --- a/emalign/align_z/utils.py +++ b/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 \ No newline at end of file + return np.abs(global_offset), ref_bboxes diff --git a/emalign/scripts/fuse_stacks_xy.py b/emalign/scripts/fuse_stacks_xy.py index ce77033..2d6fc1e 100644 --- a/emalign/scripts/fuse_stacks_xy.py +++ b/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) \ No newline at end of file + recompute_configs=args.recompute_configs, + wipe_progress_stack=args.wipe_progress_stack)