2323import json
2424import logging
2525import os
26- from typing import Any , Union
26+ from typing import Any , Optional
2727
2828import anndata as ad
2929import dask .array as da
@@ -83,7 +83,7 @@ def segment_ROI(
8383 foreground_class : int = 0 ,
8484 threshold : float = 0.5 ,
8585 min_size : int = 15 ,
86- label_dtype : Union [np .dtype , None ] = None ,
86+ label_dtype : Optional [np .dtype ] = None ,
8787 relabeling : bool = True ,
8888) -> np .ndarray :
8989 """Run the Ilastik model on a single ROI.
@@ -189,8 +189,8 @@ def ilastik_pixel_classification_segmentation(
189189 default_factory = IlastikChannel2InputModel
190190 ),
191191 input_ROI_table : str = "FOV_ROI_table" ,
192- output_ROI_table : Union [str , None ] = None ,
193- output_label_name : Union [str , None ] = None ,
192+ output_ROI_table : Optional [str ] = None ,
193+ output_label_name : Optional [str ] = None ,
194194 use_masks : bool = True ,
195195 # Ilastik-related arguments
196196 ilastik_model : str ,
@@ -503,6 +503,24 @@ def ilastik_pixel_classification_segmentation(
503503 preprocessing_kwargs = preprocessing_kwargs ,
504504 )
505505
506+ # Make 3D in case of 2D data
507+ # Check that the shape of the new label image matches the expected shape
508+ expected_shape = (
509+ e_z - s_z ,
510+ e_y - s_y ,
511+ e_x - s_x ,
512+ )
513+ logger .info (f"Expected shape: { expected_shape } " )
514+ if new_label_img .shape != expected_shape :
515+ try :
516+ new_label_img = np .broadcast_to (new_label_img , expected_shape )
517+ except Exception as err :
518+ raise ValueError (
519+ f"Shape mismatch: { new_label_img .shape } != { expected_shape } "
520+ "Between the segmented label image and expected shape in "
521+ "the zarr array."
522+ ) from err
523+
506524 if output_ROI_table :
507525 bbox_df = array_to_bounding_box_table (
508526 new_label_img ,
@@ -519,21 +537,6 @@ def ilastik_pixel_classification_segmentation(
519537 f"{ len (overlap_list )} bounding-box pairs overlap"
520538 )
521539
522- # Check that the shape of the new label image matches the expected shape
523- expected_shape = (
524- e_z - s_z ,
525- e_y - s_y ,
526- e_x - s_x ,
527- )
528- if new_label_img .shape != expected_shape :
529- try :
530- new_label_img = da .broadcast_to (new_label_img , expected_shape )
531- except :
532- raise ValueError (
533- f"Shape mismatch: { new_label_img .shape } != { expected_shape } "
534- "Between the segmented label image and expected shape in the zarr array."
535- )
536-
537540 # Compute and store 0-th level to disk
538541 da .array (new_label_img ).to_zarr (
539542 url = mask_zarr ,
0 commit comments