Quickfix

TPTBox/core/nii_wrapper.py

@@ -1255,13 +1255,14 @@ def filter_connected_components(self, labels: int |list[int]|None=None,min_volum
         """
         assert self.seg, "This only works on segmentations"
         arr = np_filter_connected_components(self.get_seg_array(), largest_k_components=max_count_component,label_ref=labels,connectivity=connectivity,return_original_labels=keep_label,min_volume=min_volume,max_volume=max_volume,removed_to_label=removed_to_label,)
-        #if keep_label and labels is not None:
-        #    if isinstance(labels,int):
-        #        labels = [labels]
-        #    old_labels = [i for i in self.unique() if i not in labels]
-        #    if len(old_labels) != 0:
-        #        s = self.extract_label(old_labels,keep_label=True)
-        #        nii[s != 0] = s[s!=0]
+        assert arr.shape == self.shape, f"Shape mismatch: {arr.shape} != {self.shape}"
+        if keep_label and labels is not None:
+            if isinstance(labels,int):
+                labels = [labels]
+            old_labels = [i for i in self.unique() if i not in labels]
+            if len(old_labels) != 0:
+                s = self.extract_label(old_labels,keep_label=True).get_array()
+                arr[s != 0] = s[s!=0]
         #print("filter",nii.unique())
         #assert max_count_component is None or nii.max() <= max_count_component, nii.unique()
         return self.set_array(arr, inplace=inplace)

TPTBox/core/poi.py

@@ -1066,7 +1066,7 @@ def calc_centroids_from_two_masks(
     org_shape = subreg_msk.shape
     # crop to mask to speed up the segmentation
     crop = vert_msk.compute_crop()
-    crop = subreg_msk.compute_crop(maximum_size=crop)
+    # crop = subreg_msk.compute_crop(maximum_size=crop)
     # crop = (slice(0, subreg_msk.shape[0]), slice(0, subreg_msk.shape[1]), slice(0, subreg_msk.shape[2]))
 
     vert_msk = vert_msk.apply_crop(crop)