updated processing sections functions

Author: hschryver

src/xenium_analysis_tools/process_xenium/generate_dataset_slides.py

@@ -15,28 +15,6 @@
     get_partial_dataset
 )
 from xenium_analysis_tools.process_xenium.process_spatialdata import read_xenium_slide
-
-def find_xenium_bundle(bundle_name, data_folder='/root/capsule/data'):
-    data_folder = Path(data_folder)
-    search_paths = [
-        data_folder / 'xenium_data',
-        data_folder / 'Xenium_output_pilot'
-    ]
-    search_paths = [path for path in search_paths if path.exists()]
-    all_dirs = np.concatenate([list(folder.iterdir()) for folder in search_paths])
-    output_folders = np.concatenate([list(folder.glob('output-*')) for folder in search_paths])
-    subfolders = np.setdiff1d(all_dirs, output_folders)
-    path_to_bundle = None
-    found_dirs = [dir for dir in output_folders if dir.name == bundle_name]
-    if found_dirs:
-        path_to_bundle = found_dirs[0]
-    else:
-        for sub in subfolders:
-            found_dirs = [dir for dir in list(sub.iterdir()) if dir.name == bundle_name]
-            if found_dirs:
-                path_to_bundle = found_dirs[0]
-                break
-    return path_to_bundle
 
 def generate_slides(dataset_name: str, config_path: str=None, select_sections: list[int]|None = None):
     """

src/xenium_analysis_tools/process_xenium/process_dataset_slides.py

@@ -1,10 +1,18 @@
+from pathlib import Path
+import tqdm
+import pandas as pd
+import numpy as np
+import spatialdata as sd
+
 from xenium_analysis_tools.utils.io_utils import (
     atomic_write_sdata, 
     is_complete, 
     is_complete_store, 
     load_config, 
     setup_logging,
-    get_sections_df
+    get_sections_df,
+    get_partial_dataset,
+    find_xenium_bundle
 )
 from xenium_analysis_tools.process_xenium.process_spatialdata import (
     process_metadata, 
@@ -18,11 +26,6 @@
 from xenium_analysis_tools.process_xenium.validate_sections import (
     plot_section_bboxes
 )
-from pathlib import Path
-import tqdm
-import pandas as pd
-import numpy as np
-import spatialdata as sd
 
 def process_slides(dataset_name: str, config_path: str=None, select_sections: list[int]|None = None, slides_parent_folder='data'):
     """
@@ -31,63 +34,69 @@ def process_slides(dataset_name: str, config_path: str=None, select_sections: li
 
     # ---- Set up ----
     config = load_config(config_path)
+
+    # Paths/directories
     paths = config['paths']
     processing_config = config['processing_control']
     raw_data_folder = Path(paths['data_root']) / dataset_name
     slide_sd_path = Path(paths[f'{slides_parent_folder}_root']) / f"{dataset_name}{processing_config['save_initial_dataset_suffix']}"
     save_sections_parent_folder = processing_config['save_processed_data_parent_folder']
     save_sections_path = Path(paths[f'{save_sections_parent_folder}_root']) / f"{dataset_name}{processing_config['save_processed_dataset_suffix']}"
     save_sections_path.mkdir(parents=True, exist_ok=True)
+
+    # Logger
     logger, log_file_path = setup_logging(save_sections_path)
+
+    # Print out where sections are being saved
     logger.info(f"Dataset Name: {dataset_name}")
     logger.info(f"Configuration loaded from {config_path}")
     logger.info(f"Raw data folder: {raw_data_folder}")
     logger.info(f"Slides are being loaded from: {slide_sd_path}")
     logger.info(f"Processed sections will be saved to: {save_sections_path}")
+
+    # If specified, copy sections from data folder instead of re-generating
+    if processing_config['check_data_folder_slides']:
+        logger.info("Checking and copying slides from data folder if exist...")
+        data_folder_slides_path = Path(paths['data_root']) / f'{dataset_name}{processing_config["save_processed_dataset_suffix"]}'
+        get_partial_dataset(data_folder_slides_path, save_sections_path, pattern='slide_*', subset_ids=select_sections)
+
+    # Get the slides information
     sections_df = get_sections_df(raw_data_folder)
+
     # Limit sections, if specified
     if select_sections is not None:
         logger.info(f"Limiting processing to sections: {select_sections}")
         sections_df = sections_df[sections_df['section'].isin(select_sections)]
+    
+    # Set up processing loop
     logger.info(f"Total slides found: {len(sections_df)}")
     unique_slides = sections_df.groupby('slide_id')
 
-    # ---- (Optional) Move data from a data asset to results folder ----
-    # Load partially processed sections from data asset and save to results to further process
-    if config['processing_control'].get('load_processed_from_asset', False):
-        logger.info(f"Looking for processed sections in a data asset...")
-        dataset_asset_folder = Path(config['paths']['data_root']) / f"{dataset_name}{config['processing_control']['save_processed_dataset_suffix']}"
-        if dataset_asset_folder.exists():
-            logger.info(f"Loading processed sections from data asset at {dataset_asset_folder}")
-            sections_in_folder = list(dataset_asset_folder.glob('section_*.zarr'))
-            for section_zarr_path in tqdm.tqdm(sections_in_folder, desc="Moving processed sections from asset"):
-                section_save_path = save_sections_path / section_zarr_path.name
-                if section_save_path.exists() and is_complete_store(section_save_path):
-                    logger.info(f"Section {section_zarr_path.name} already exists in results folder. Skipping.")
-                    continue
-                logger.info(f"Loading section {section_zarr_path.stem} from asset and saving to results folder...")
-                try:
-                    sdata = sd.read_zarr(section_zarr_path)
-                    atomic_write_sdata(sdata, section_save_path, overwrite=True)
-                    del sdata
-                except Exception as e:
-                    logger.error(f"Error loading section {section_zarr_path.name} from asset: {e}")
-                    continue
-        else:
-            logger.info(f"No processed data asset found at {dataset_asset_folder}. Continuing without loading from asset.")
-
     # ---- Run processing ----
-    logger.info(f"Starting processing for dataset: {dataset_name}")
-    unique_slides = sections_df.groupby('slide_id')
-    for slide_id in tqdm.tqdm(unique_slides.groups.keys(), desc="Processing slides", unit="slide", total=len(unique_slides.groups.keys())):
+    for slide_id in tqdm.tqdm(unique_slides.groups.keys(), 
+                                desc="Processing slides", 
+                                unit="slide", 
+                                total=len(unique_slides.groups.keys())):
+
+        # Get slide information
         group = unique_slides.get_group(slide_id)
         slide_row = group.iloc[0]
         raw_slide_path = raw_data_folder / slide_row['dir']
+
+        # Make sure Xenium bundle is valid, or find alternative location
+        if not (raw_slide_path / 'experiment.xenium').exists():
+            xenium_bundle_path = find_xenium_bundle(raw_slide_path.name, data_folder=paths['data_root'])
+            if xenium_bundle_path is not None:
+                logger.info(f"experiment.xenium not found in {raw_slide_path}. Using found bundle at {xenium_bundle_path}")
+                raw_slide_path = xenium_bundle_path
+
         slide_sdata_path = slide_sd_path / f"{processing_config['save_initial_dataset_prefix']}{slide_id}.zarr"
         if not slide_sdata_path.exists():
             logger.warning(f"Slide data not found for slide {slide_id} at {slide_sdata_path}! Skipping.")
             continue
         slide_sections = slide_row['slide_sections']
+
+        # Check which sections need processing
         process_sections = []
         for section in slide_sections:
             section_zarr = f"{processing_config['save_processed_dataset_prefix']}{section}.zarr"
@@ -99,10 +108,10 @@ def process_slides(dataset_name: str, config_path: str=None, select_sections: li
         if not process_sections:
             logger.info(f"All sections for slide {slide_id} are already processed. Skipping slide.")
             continue
-        logger.info(f"Processing slide {slide_id} sections: {[str(s) for s in process_sections]}")
 
         # --- Process slide SpatialData ---
         # Load slide SpatialData
+        logger.info(f"Processing slide {slide_id} sections: {[str(s) for s in process_sections]}")
         slide_sdata = sd.read_zarr(slide_sdata_path)
 
         # Get additional metadata from raw data xenium bundle

src/xenium_analysis_tools/utils/io_utils.py

@@ -143,6 +143,28 @@ def safe_copy_tree(src: Path, dst: Path):
 
     shutil.copytree(src, dst)
 
+def find_xenium_bundle(bundle_name, data_folder='/root/capsule/data'):
+    data_folder = Path(data_folder)
+    search_paths = [
+        data_folder / 'xenium_data',
+        data_folder / 'Xenium_output_pilot'
+    ]
+    search_paths = [path for path in search_paths if path.exists()]
+    all_dirs = np.concatenate([list(folder.iterdir()) for folder in search_paths])
+    output_folders = np.concatenate([list(folder.glob('output-*')) for folder in search_paths])
+    subfolders = np.setdiff1d(all_dirs, output_folders)
+    path_to_bundle = None
+    found_dirs = [dir for dir in output_folders if dir.name == bundle_name]
+    if found_dirs:
+        path_to_bundle = found_dirs[0]
+    else:
+        for sub in subfolders:
+            found_dirs = [dir for dir in list(sub.iterdir()) if dir.name == bundle_name]
+            if found_dirs:
+                path_to_bundle = found_dirs[0]
+                break
+    return path_to_bundle
+
 def get_partial_dataset(source_path, dest_path, pattern='section_*', subset_ids=None):
     """Copy slide data from source to destination, handling incomplete files."""
     # Find matches