added points df to zstacks

Author: hschryver

src/xenium_analysis_tools/alignment/generate_images.py

@@ -1,6 +1,6 @@
 import spatialdata as sd
 from xenium_analysis_tools.utils.sd_utils import add_micron_coord_sys
-from spatialdata.models import Image3DModel, Labels3DModel
+from spatialdata.models import Image3DModel, Labels3DModel, PointsModel
 from pathlib import Path
 import pandas as pd
 import numpy as np
@@ -63,9 +63,9 @@ def get_zstacks_dict(zstacks_folder, channels=['gcamp', 'dextran']):
 
     for stack_ind, stack_folder in enumerate(stack_dirs):
         stack_info = {
-            'stack_name': stack_folder.name,
-            'stack_size': _extract_stack_size(stack_folder.name),
-            'stack_channels': [ch for ch in channels if ch in stack_folder.name.lower()],
+            'zstack_name': stack_folder.name,
+            'zstack_size': _extract_zstack_size(stack_folder.name),
+            'zstack_channels': [ch for ch in channels if ch in stack_folder.name.lower()],
             'metadata_jsons': {'registration': None, 'roi_groups': None, 'scanimage': None},
             'channel_tifs': {}
         }
@@ -92,9 +92,9 @@ def get_zstacks_dict(zstacks_folder, channels=['gcamp', 'dextran']):
 
     return zstacks_dict
 
-def _extract_stack_size(stack_name):
+def _extract_zstack_size(zstack_name):
     """Extract width x height x depth from stack name."""
-    size_pattern = re.search(r'(\d+)x(\d+)x(\d+)', stack_name)
+    size_pattern = re.search(r'(\d+)x(\d+)x(\d+)', zstack_name)
     if size_pattern:
         width, height, depth = map(int, size_pattern.groups())
         return {"width": width, "height": height, "depth": depth}
@@ -110,7 +110,7 @@ def _categorize_json_file(filename_lower):
         return 'scanimage'
     return None
 
-def get_zstack(zstacks_dict, zstack_ind=None, zstack_name=None, stack_size=None, channels=None):    
+def get_zstack(zstacks_dict, zstack_ind=None, zstack_name=None, zstack_size=None, zstack_channels=None):    
     if zstack_ind is not None:
         if zstack_ind not in zstacks_dict:
             raise ValueError(f"Z-stack index {zstack_ind} not found in zstacks_dict.")
@@ -127,37 +127,36 @@ def _find_matches(criterion_func, criterion_name, criterion_value):
             return zstacks_dict[matches[0]]
 
         # Handle multiple matches with optional channel filtering
-        if channels is not None:
+        if zstack_channels is not None:
             channel_matches = [
                 i for i in matches 
-                if set(zstacks_dict[i]['stack_channels']) == set(channels)
+                if set(zstacks_dict[i]['zstack_channels']) == set(zstack_channels)
             ]
             if len(channel_matches) == 1:
                 return zstacks_dict[channel_matches[0]]
             elif len(channel_matches) > 1:
-                raise ValueError(f"Multiple z-stacks found with {criterion_name} {criterion_value} and channels {channels}. Found {len(channel_matches)} matches.")
+                raise ValueError(f"Multiple z-stacks found with {criterion_name} {criterion_value} and channels {zstack_channels}. Found {len(channel_matches)} matches.")
             else:
-                raise ValueError(f"No z-stack found with {criterion_name} {criterion_value} and channels {channels}.")
-        
+                raise ValueError(f"No z-stack found with {criterion_name} {criterion_value} and channels {zstack_channels}.")
         raise ValueError(f"Multiple z-stacks found with {criterion_name} {criterion_value}. Found {len(matches)} matches. Consider specifying channels parameter.")
 
     if zstack_name is not None:
         return _find_matches(
-            lambda stack: stack['stack_name'] == zstack_name,
+            lambda stack: stack['zstack_name'] == zstack_name,
             "Z-stack name", zstack_name
         )
 
-    if stack_size is not None:
+    if zstack_size is not None:
         return _find_matches(
             lambda stack: (
-                stack['stack_size']['width'] == stack_size['width'] and
-                stack['stack_size']['height'] == stack_size['height'] and
-                stack['stack_size']['depth'] == stack_size['depth']
+                stack['zstack_size']['width'] == zstack_size['width'] and
+                stack['zstack_size']['height'] == zstack_size['height'] and
+                stack['zstack_size']['depth'] == zstack_size['depth']
             ),
-            "Stack size", stack_size
+            "Stack size", zstack_size
         )
 
-    raise ValueError("Either zstack_ind, zstack_name, or stack_size must be provided.")
+    raise ValueError("Either zstack_ind, zstack_name, or zstack_size must be provided.")
 
 def get_alignment_data_paths(dataset_id, 
                             data_root=Path('/root/capsule/data'),
@@ -182,68 +181,84 @@ def get_alignment_data_paths(dataset_id,
 
     return paths
 
-def get_zstack_sdata(stack, zstack_masks=None, use_shared_coords=True):
+def get_label_params(label_obj, id_name='cell'):
+    from skimage.measure import regionprops
+    labels = label_obj.values
+    props = regionprops(labels)
+    data = [
+        {f'{id_name}_id': p.label, 
+        'z': p.centroid[0], 
+        'y': p.centroid[1], 
+        'x': p.centroid[2],
+        'area': p.area,
+        'bbox': p.bbox}
+        for p in props
+    ]
+    df = pd.DataFrame(data)
+    return df
+
+def get_zstack_sdata(stack, zstack_masks=None, get_centroids_as_points=True):
     # Create the z-stack image array
-    num_channels = len(stack['stack_channels'])
+    num_channels = len(stack['zstack_channels'])
     chans = []
     if num_channels > 1:
         for ch_ind in range(num_channels):
             chan_array = create_zstack_array(tif_path=stack['channel_tifs'][ch_ind]['chan_tif_path'], 
-                    fov_x_um=stack['stack_size']['width'], 
-                    fov_y_um=stack['stack_size']['height'], 
-                    fov_z_um=stack['stack_size']['depth'])
+                    fov_x_um=stack['zstack_size']['width'], 
+                    fov_y_um=stack['zstack_size']['height'], 
+                    fov_z_um=stack['zstack_size']['depth'])
             chans.append(chan_array)
         zstack_img = xr.concat(chans, dim='c')
-        zstack_img['c'] = stack['stack_channels']
+        zstack_img['c'] = stack['zstack_channels']
     else:
         zstack_img = create_zstack_array(tif_path=stack['channel_tifs'][0]['chan_tif_path'], 
-                    fov_x_um=stack['stack_size']['width'], 
-                    fov_y_um=stack['stack_size']['height'], 
-                    fov_z_um=stack['stack_size']['depth'])
-        zstack_img['c'] = stack['stack_channels']
-
-    if use_shared_coords:   
-        reg_json_path = stack['metadata_jsons']['registration']
-        with open(reg_json_path) as f:
-            reg_json = json.load(f) 
-        if 'z_steps' in reg_json.keys() and len(reg_json['z_steps'])==zstack_img.sizes['z']:
-            print("Using shared z coordinates for images")
-            zstack_img.coords['z'] = reg_json['z_steps']
+                    fov_x_um=stack['zstack_size']['width'], 
+                    fov_y_um=stack['zstack_size']['height'], 
+                    fov_z_um=stack['zstack_size']['depth'])
+        zstack_img['c'] = stack['zstack_channels']
 
     # Parse into Image3DModel
     zstack_img = Image3DModel.parse(
                 zstack_img,
                 dims=['c', 'z', 'y', 'x'],
-                c_coords=stack['stack_channels'],
+                c_coords=stack['zstack_channels'],
                 chunks='auto',
             )
 
-    # Make the SpatialData object
-    zstack_sdata = sd.SpatialData(
-            images={'zstack': zstack_img},
-        )
-
     if zstack_masks is not None:
+        zstack_labels = {}
         # Get labels for each channel
         for mask_ind, masks in zstack_masks['channel_tifs'].items():
-            channel_name = zstack_masks['stack_channels'][mask_ind]
+            channel_name = zstack_masks['zstack_channels'][mask_ind]
             zstack_label = create_zstack_array(tif_path=masks['chan_tif_path'], 
-                        fov_x_um=zstack_masks['stack_size']['width'], 
-                        fov_y_um=zstack_masks['stack_size']['height'], 
-                        fov_z_um=zstack_masks['stack_size']['depth'],
+                        fov_x_um=zstack_masks['zstack_size']['width'], 
+                        fov_y_um=zstack_masks['zstack_size']['height'], 
+                        fov_z_um=zstack_masks['zstack_size']['depth'],
                         add_chan=False)
 
-            if use_shared_coords:   
-                if 'z_steps' in reg_json.keys() and len(reg_json['z_steps'])==zstack_label.sizes['z']:
-                    print("Using shared z coordinates for labels")
-                    zstack_label.coords['z'] = reg_json['z_steps']
-
             zstack_label = Labels3DModel.parse(
                         zstack_label,
                         dims=['z', 'y', 'x'],
                         chunks='auto',
                     )
-            zstack_sdata.labels[f"{channel_name}_labels"] = zstack_label
+            zstack_labels[f"{channel_name}_labels"] = zstack_label
+
+        if get_centroids_as_points:
+            zstack_points = {}
+            # Get label parameters  add as points
+            for label_name, labels_obj in zstack_labels.items():
+                chan_name = label_name.replace('_labels','')
+                cells_df = get_label_params(labels_obj, id_name=chan_name)
+                print(f"# {chan_name} segmented cells: {len(cells_df)}")
+                cells_df = PointsModel.parse(cells_df)
+                zstack_points[f"{chan_name}_cells"] = cells_df
+
+    # Assemble SpatialData
+    zstack_sdata = sd.SpatialData(
+        images={'zstack': zstack_img},
+        labels={**zstack_labels} if zstack_masks is not None else {},
+        points={**zstack_points} if (zstack_masks is not None and get_centroids_as_points) else {}
+    )    
 
     # Determine pixel sizes
     if zstack_sdata['zstack'].attrs['pixel_size_um_x'] == zstack_sdata['zstack'].attrs['pixel_size_um_y']: