Merge pull request #75 from MannLabs/fix_selection

fix segmentation shape problem in the selection pipeline

Author: sophiamaedler

src/scportrait/pipeline/selection.py

@@ -170,7 +170,6 @@ def process(
         )
 
         segmentation = tempmmap.mmap_array_from_path(self.path_seg_mask)
-        print(segmentation)
 
         # create segmentation loader
         sl = SegmentationLoader(
@@ -179,6 +178,13 @@ def process(
             processes=self.config["processes_cell_sets"],
         )
 
+        if len(segmentation.shape) == 3:
+            segmentation = np.squeeze(segmentation)
+        else:
+            raise ValueError(
+                f"Segmentation shape is not correct. Expected 2D array, got {segmentation.shape}"
+                )
+
         # get shape collections
         shape_collection = sl(segmentation, self.cell_sets, self.calibration_marker)