From 449ebcd87f5e631505b94500d0dc98fa0d2fe30c Mon Sep 17 00:00:00 2001 From: gu53lif Date: Wed, 30 Sep 2026 10:19:46 +0200 Subject: [PATCH 1/3] bug fix _rib_assingn --- TPTBox/segmentation/rib/_rib_assign.py | 133 +++++++++++++++++++------ 1 file changed, 104 insertions(+), 29 deletions(-) diff --git a/TPTBox/segmentation/rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py index 426880d..39025bc 100644 --- a/TPTBox/segmentation/rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -47,7 +47,9 @@ def _thoracic_like_labels(labels: Iterable[int]) -> list[int]: return [*sorted(v for v in labels if 7 <= int(v) <= 20), 28] -def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> tuple[list[RibCandidate], list[RibCandidate]]: +def _split_ccs_by_side( + rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI +) -> tuple[list[RibCandidate], list[RibCandidate]]: right_axis = cms_vert.get_axis("R") inf_axis = cms_vert.get_axis("I") # ref_x = cms_vert[ref_vert, 50][right_axis] @@ -59,7 +61,12 @@ def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> continue center = cms_cc[cc, 50] center2 = cms_cc2[cc, 50] if (cc, 50) in cms_cc2 else center # noqa: SIM401 - cand = RibCandidate(cc_label=cc, volume=vols.get(cc, 0), z=center[inf_axis], x=center2[right_axis]) + cand = RibCandidate( + cc_label=cc, + volume=vols.get(cc, 0), + z=center[inf_axis], + x=center2[right_axis], + ) distances = cms_vert.calculate_distances_cord(center2) min_key = min(distances, key=distances.get) # type: ignore # x < vertebra center => patient right in RAS-like orientation @@ -74,14 +81,20 @@ def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> return left, right -def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None) -> float: +def _touching_surface( + mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None +) -> float: """Return weighted number of voxel faces shared by two masks. Touching along ``up_down_axis`` is weighted by ``up_down_weight``. All other axes have weight 1.0. """ if axis_weights is None: - axis_weights = {nii.get_axis("S"): 0.01, nii.get_axis("A"): 0.1, nii.get_axis("R"): 1} + axis_weights = { + nii.get_axis("S"): 0.01, + nii.get_axis("A"): 0.1, + nii.get_axis("R"): 1, + } surface = 0.0 for axis in range(mask_a.ndim): @@ -92,7 +105,8 @@ def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_wei sl2[axis] = slice(1, None) touching = ( - np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() + np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() + + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() ) weight = axis_weights.get(axis, 1) @@ -102,7 +116,12 @@ def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_wei def _try_erosion_split( - cc_label: int, binary_cc: NII, erosion_pixels: int, min_volume: int, _pass=0, verbose=True + cc_label: int, + binary_cc: NII, + erosion_pixels: int, + min_volume: int, + _pass=0, + verbose=True, ) -> tuple[int, list[np.ndarray]] | None: """Try to split one CC by erosion. @@ -123,9 +142,13 @@ def _try_erosion_split( elif _pass == 1: eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) elif _pass == 2: - eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="R") + eroded = cc_nii.erode_msk( + n_pixel=erosion_pixels, verbose=False, ignore_direction="R" + ) elif _pass == 3: - eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="A") + eroded = cc_nii.erode_msk( + n_pixel=erosion_pixels, verbose=False, ignore_direction="A" + ) else: eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) except Exception as e: @@ -144,7 +167,9 @@ def _try_erosion_split( for sub in sub_labels.copy(): count = int((infected_arr == sub).sum()) - should_merge = count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 + should_merge = ( + count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 + ) if not should_merge: continue # print("merge", count, full_size, count / full_size) @@ -157,7 +182,12 @@ def _try_erosion_split( target = remaining[0] # print("merge", target) else: - target = max(remaining, key=lambda candidate: _touching_surface(sub_mask, infected_arr == candidate, infected)) + target = max( + remaining, + key=lambda candidate: _touching_surface( + sub_mask, infected_arr == candidate, infected + ), + ) # print("merge of many", target) infected_arr[infected_arr == sub] = target sub_labels.remove(sub) @@ -219,7 +249,9 @@ def split_touching_rib_ccs( next_label = int(rib_cc.max()) + 1 if rib_cc.shape and rib_cc.max() > 0 else 1 u = {int(x) for x in rib_cc.unique() if x != 0} # ribs counte twice (left/right) - expected_number_of_ccs = len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 + expected_number_of_ccs = ( + len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 + ) if expected_number_of_ccs == 0: expected_number_of_ccs = 24 @@ -231,8 +263,12 @@ def split_touching_rib_ccs( binary_ccs = [(cc_label, rib_cc.extract_label(cc_label)) for cc_label in labels] if num_workers == 1: results = [] - for cc_label, binary_cc in tqdm(binary_ccs, total=len(binary_ccs), desc="Separate RIBs"): - result = _try_erosion_split(cc_label, binary_cc, erosion_pixels, min_volume, _pass) + for cc_label, binary_cc in tqdm( + binary_ccs, total=len(binary_ccs), desc="Separate RIBs" + ): + result = _try_erosion_split( + cc_label, binary_cc, erosion_pixels, min_volume, _pass + ) if result is not None: results.append(result) else: @@ -240,12 +276,21 @@ def split_touching_rib_ccs( with ThreadPoolExecutor(max_workers=num_workers) as executor: futures = [ - executor.submit(_try_erosion_split, cc_label, binary_cc, erosion_pixels, min_volume, _pass) + executor.submit( + _try_erosion_split, + cc_label, + binary_cc, + erosion_pixels, + min_volume, + _pass, + ) for cc_label, binary_cc in binary_ccs ] # tqdm advances immediately whenever an individual future completes. - for future in tqdm(as_completed(futures), total=len(futures), desc="Separate RIBs"): + for future in tqdm( + as_completed(futures), total=len(futures), desc="Separate RIBs" + ): result = future.result() if result is not None: results.append(result) @@ -268,7 +313,9 @@ def split_touching_rib_ccs( u.add(new_lbl) if verbose: - logger.print(f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})") + logger.print( + f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})" + ) if short_cut and rib_cc.max() >= expected_number_of_ccs: break @@ -338,18 +385,29 @@ def assign_ribs_to_vert_segmentation( vert_seg = vert_seg.reorient() sem_seg = sem_seg.reorient() rib_seg = rib_seg.reorient(verbose=verbose) - vert_pred = vert_seg.extract_label([Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1]) + vert_pred = vert_seg.extract_label( + [Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1] + ) rib_seg = rib_seg.extract_label([left_id, right_id], keep_label=True) # type: ignore logger.on_debug(f"{rib_seg.unique()=}") # Remove rib voxels overlapping vertebrae rib_seg[vert_pred != 0] = 0 - rib_cc = rib_seg.filter_connected_components(None, min_volume=min_volume, keep_label=False) + rib_cc = rib_seg.filter_connected_components( + None, min_volume=min_volume, keep_label=False + ) cms_vert = calc_centroids(vert_seg) if split_touching: rib_cc = split_touching_rib_ccs( - rib_cc, erosion_pixels=erosion_pixels, min_volume=min_volume, vert_ids=vert_ids, short_cut=short_cut, verbose=verbose + rib_cc, + erosion_pixels=erosion_pixels, + min_volume=min_volume, + vert_ids=vert_ids, + short_cut=short_cut, + verbose=verbose, ) - cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) + cms_cc = calc_centroids( + rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5) + ) cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) vert_labels = _thoracic_like_labels(vert_seg.unique()) @@ -373,12 +431,17 @@ def assign_ribs_to_vert_segmentation( assert isinstance(vert_labels, (tuple, list)), type(vert_labels) if vert_labels and ( - (vert_labels[0] != 28 and vert_labels[1] != 28) or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) + (vert_labels[0] != 28 and vert_labels[1] != 28) + or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) ): inf_axis = cms_vert.get_axis("I") t1_z = cms_vert[vert_labels[0], 50][inf_axis] # use median vertebral spacing as adaptive threshold - spacing = abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) if len(vert_labels) > 1 else 30 + spacing = ( + abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) + if len(vert_labels) > 1 + else 30 + ) threshold = max(5, spacing * 0.6) nearest_left = abs(left_ccs[0].z - t1_z) if left_ccs else float("inf") @@ -405,20 +468,32 @@ def assign_ribs_to_vert_segmentation( logger.on_debug(f"{rib_inst.unique()=}") # Merge rib assignments back into the original vert / sem segmentations # without disturbing existing (non-rib) labels. Skip unmatched CCs (sentinel error_value=255). - matched = (rib_inst != 0) & (rib_inst != error_value) + print("rib_inst", rib_inst.dtype, rib_inst.seg) + matched = np.logical_and(rib_inst != 0, rib_inst != error_value) vert_seg[matched] = rib_inst[matched] - sem_seg[rib_seg != 0] = rib_seg.map_labels({left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value})[rib_seg != 0] # type: ignore + sem_seg[rib_seg != 0] = rib_seg.map_labels( + {left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value} + )[rib_seg != 0] # type: ignore undefined = sum(v == error_value for v in rib_vert_map.values()) logger.print(f"Unmatched rib CCs: {undefined}") if add_error: # or split_touching: vert_seg[rib_inst == error_value] = error_value if np.any(vert_seg.get_seg_array() == error_value): - vert_seg2 = vert_seg.remove_labels(error_value).infect(vert_seg.extract_label(error_value), verbose=False) + vert_seg2 = vert_seg.remove_labels(error_value).infect( + vert_seg.extract_label(error_value), verbose=False + ) vert_seg[vert_seg != vert_seg2] = vert_seg2[vert_seg != vert_seg2] - vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = error_value - vert_seg[np.logical_and(vert_seg == 0, sem_seg.extract_label([Location.Rib_Left.value, Location.Rib_Right.value] == 1))] = ( - error_value - ) + vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = ( + error_value + ) + vert_seg[ + np.logical_and( + vert_seg == 0, + sem_seg.extract_label( + [Location.Rib_Left.value, Location.Rib_Right.value] + ) == 1, + ) + ] = error_value return vert_seg.reorient_(ori), sem_seg.reorient_(ori) From 41e4a15aa7e9ede360e24813e1e8a6c6aed68253 Mon Sep 17 00:00:00 2001 From: TanjaLe <105376287+TanjaLe@users.noreply.github.com> Date: Wed, 30 Sep 2026 08:21:34 +0000 Subject: [PATCH 2/3] style fixes by ruff --- TPTBox/segmentation/rib/_rib_assign.py | 84 +++++++------------------- 1 file changed, 21 insertions(+), 63 deletions(-) diff --git a/TPTBox/segmentation/rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py index 39025bc..44302c9 100644 --- a/TPTBox/segmentation/rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -47,9 +47,7 @@ def _thoracic_like_labels(labels: Iterable[int]) -> list[int]: return [*sorted(v for v in labels if 7 <= int(v) <= 20), 28] -def _split_ccs_by_side( - rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI -) -> tuple[list[RibCandidate], list[RibCandidate]]: +def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> tuple[list[RibCandidate], list[RibCandidate]]: right_axis = cms_vert.get_axis("R") inf_axis = cms_vert.get_axis("I") # ref_x = cms_vert[ref_vert, 50][right_axis] @@ -81,9 +79,7 @@ def _split_ccs_by_side( return left, right -def _touching_surface( - mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None -) -> float: +def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None) -> float: """Return weighted number of voxel faces shared by two masks. Touching along ``up_down_axis`` is weighted by ``up_down_weight``. @@ -105,8 +101,7 @@ def _touching_surface( sl2[axis] = slice(1, None) touching = ( - np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() - + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() + np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() ) weight = axis_weights.get(axis, 1) @@ -142,13 +137,9 @@ def _try_erosion_split( elif _pass == 1: eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) elif _pass == 2: - eroded = cc_nii.erode_msk( - n_pixel=erosion_pixels, verbose=False, ignore_direction="R" - ) + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="R") elif _pass == 3: - eroded = cc_nii.erode_msk( - n_pixel=erosion_pixels, verbose=False, ignore_direction="A" - ) + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="A") else: eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) except Exception as e: @@ -167,9 +158,7 @@ def _try_erosion_split( for sub in sub_labels.copy(): count = int((infected_arr == sub).sum()) - should_merge = ( - count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 - ) + should_merge = count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 if not should_merge: continue # print("merge", count, full_size, count / full_size) @@ -184,9 +173,7 @@ def _try_erosion_split( else: target = max( remaining, - key=lambda candidate: _touching_surface( - sub_mask, infected_arr == candidate, infected - ), + key=lambda candidate: _touching_surface(sub_mask, infected_arr == candidate, infected), ) # print("merge of many", target) infected_arr[infected_arr == sub] = target @@ -249,9 +236,7 @@ def split_touching_rib_ccs( next_label = int(rib_cc.max()) + 1 if rib_cc.shape and rib_cc.max() > 0 else 1 u = {int(x) for x in rib_cc.unique() if x != 0} # ribs counte twice (left/right) - expected_number_of_ccs = ( - len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 - ) + expected_number_of_ccs = len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 if expected_number_of_ccs == 0: expected_number_of_ccs = 24 @@ -263,12 +248,8 @@ def split_touching_rib_ccs( binary_ccs = [(cc_label, rib_cc.extract_label(cc_label)) for cc_label in labels] if num_workers == 1: results = [] - for cc_label, binary_cc in tqdm( - binary_ccs, total=len(binary_ccs), desc="Separate RIBs" - ): - result = _try_erosion_split( - cc_label, binary_cc, erosion_pixels, min_volume, _pass - ) + for cc_label, binary_cc in tqdm(binary_ccs, total=len(binary_ccs), desc="Separate RIBs"): + result = _try_erosion_split(cc_label, binary_cc, erosion_pixels, min_volume, _pass) if result is not None: results.append(result) else: @@ -288,9 +269,7 @@ def split_touching_rib_ccs( ] # tqdm advances immediately whenever an individual future completes. - for future in tqdm( - as_completed(futures), total=len(futures), desc="Separate RIBs" - ): + for future in tqdm(as_completed(futures), total=len(futures), desc="Separate RIBs"): result = future.result() if result is not None: results.append(result) @@ -313,9 +292,7 @@ def split_touching_rib_ccs( u.add(new_lbl) if verbose: - logger.print( - f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})" - ) + logger.print(f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})") if short_cut and rib_cc.max() >= expected_number_of_ccs: break @@ -385,16 +362,12 @@ def assign_ribs_to_vert_segmentation( vert_seg = vert_seg.reorient() sem_seg = sem_seg.reorient() rib_seg = rib_seg.reorient(verbose=verbose) - vert_pred = vert_seg.extract_label( - [Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1] - ) + vert_pred = vert_seg.extract_label([Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1]) rib_seg = rib_seg.extract_label([left_id, right_id], keep_label=True) # type: ignore logger.on_debug(f"{rib_seg.unique()=}") # Remove rib voxels overlapping vertebrae rib_seg[vert_pred != 0] = 0 - rib_cc = rib_seg.filter_connected_components( - None, min_volume=min_volume, keep_label=False - ) + rib_cc = rib_seg.filter_connected_components(None, min_volume=min_volume, keep_label=False) cms_vert = calc_centroids(vert_seg) if split_touching: rib_cc = split_touching_rib_ccs( @@ -405,9 +378,7 @@ def assign_ribs_to_vert_segmentation( short_cut=short_cut, verbose=verbose, ) - cms_cc = calc_centroids( - rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5) - ) + cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) vert_labels = _thoracic_like_labels(vert_seg.unique()) @@ -431,17 +402,12 @@ def assign_ribs_to_vert_segmentation( assert isinstance(vert_labels, (tuple, list)), type(vert_labels) if vert_labels and ( - (vert_labels[0] != 28 and vert_labels[1] != 28) - or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) + (vert_labels[0] != 28 and vert_labels[1] != 28) or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) ): inf_axis = cms_vert.get_axis("I") t1_z = cms_vert[vert_labels[0], 50][inf_axis] # use median vertebral spacing as adaptive threshold - spacing = ( - abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) - if len(vert_labels) > 1 - else 30 - ) + spacing = abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) if len(vert_labels) > 1 else 30 threshold = max(5, spacing * 0.6) nearest_left = abs(left_ccs[0].z - t1_z) if left_ccs else float("inf") @@ -472,28 +438,20 @@ def assign_ribs_to_vert_segmentation( matched = np.logical_and(rib_inst != 0, rib_inst != error_value) vert_seg[matched] = rib_inst[matched] - sem_seg[rib_seg != 0] = rib_seg.map_labels( - {left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value} - )[rib_seg != 0] # type: ignore + sem_seg[rib_seg != 0] = rib_seg.map_labels({left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value})[rib_seg != 0] # type: ignore undefined = sum(v == error_value for v in rib_vert_map.values()) logger.print(f"Unmatched rib CCs: {undefined}") if add_error: # or split_touching: vert_seg[rib_inst == error_value] = error_value if np.any(vert_seg.get_seg_array() == error_value): - vert_seg2 = vert_seg.remove_labels(error_value).infect( - vert_seg.extract_label(error_value), verbose=False - ) + vert_seg2 = vert_seg.remove_labels(error_value).infect(vert_seg.extract_label(error_value), verbose=False) vert_seg[vert_seg != vert_seg2] = vert_seg2[vert_seg != vert_seg2] - vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = ( - error_value - ) + vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = error_value vert_seg[ np.logical_and( vert_seg == 0, - sem_seg.extract_label( - [Location.Rib_Left.value, Location.Rib_Right.value] - ) == 1, + sem_seg.extract_label([Location.Rib_Left.value, Location.Rib_Right.value]) == 1, ) ] = error_value return vert_seg.reorient_(ori), sem_seg.reorient_(ori) From 2716437a05c4de22b8211d6836ef2acc99fe195a Mon Sep 17 00:00:00 2001 From: gu53lif Date: Fri, 2 Oct 2026 11:23:38 +0200 Subject: [PATCH 3/3] deepali point registration and clean up --- .../_deformable/multilabel_segmentation.py | 146 ++++++++++++++---- TPTBox/segmentation/rib/_rib_assign.py | 86 ++++++++--- 2 files changed, 180 insertions(+), 52 deletions(-) diff --git a/TPTBox/registration/_deformable/multilabel_segmentation.py b/TPTBox/registration/_deformable/multilabel_segmentation.py index d4b7949..973679c 100644 --- a/TPTBox/registration/_deformable/multilabel_segmentation.py +++ b/TPTBox/registration/_deformable/multilabel_segmentation.py @@ -10,8 +10,10 @@ from TPTBox.core.poi_fun.poi_global import POI_Global from TPTBox.registration._deformable.deformable_reg import Deformable_Registration from TPTBox.registration._ridged_intensity.affine_deepali import Tether_Seg -from TPTBox.registration._ridged_points.deepali_point_registration import Deepali_Point_Registration -from TPTBox.registration._ridged_points.point_registration import Point_Registration +from TPTBox.registration._ridged_points.deepali_point_registration import ( + Deepali_Point_Registration as Point_Registration, +) +# from TPTBox.registration._ridged_points.point_registration import Point_Registration _NIIOrPOI = TypeVar("_NIIOrPOI", NII, POI) @@ -45,7 +47,9 @@ def _flip_r_axis(x: _NIIOrPOI) -> _NIIOrPOI: A new object of the same type with the R-axis reversed. """ if isinstance(x, POI_Global): - raise TypeError("_flip_r_axis: POI_Global has no shape/axis; resample to a voxel grid first.") + raise TypeError( + "_flip_r_axis: POI_Global has no shape/axis; resample to a voxel grid first." + ) axis = x.get_axis("R") if isinstance(x, NII): slicer = _r_axis_slicer(axis) @@ -171,12 +175,20 @@ def __init__( # noqa: C901 if poi_target_cms is not None: poi_target_cms = _flip_r_axis(poi_target_cms) if poi_target_cms is None: - x = target_seg.extract_label(cms_ids, keep_label=True) if cms_ids else target_seg + x = ( + target_seg.extract_label(cms_ids, keep_label=True) + if cms_ids + else target_seg + ) poi_target = calc_centroids(x, second_stage=40, bar=True) # TODO REMOVE else: poi_target = poi_target_cms.resample_from_to(target_seg) if poi_cms is None: - x = atlas_seg.extract_label(cms_ids, keep_label=True) if cms_ids else atlas_seg + x = ( + atlas_seg.extract_label(cms_ids, keep_label=True) + if cms_ids + else atlas_seg + ) poi_cms = calc_centroids(x, second_stage=40, bar=True) if not poi_cms.assert_affine(atlas_seg, raise_error=False): poi_cms = poi_cms.resample_from_to(atlas_seg) @@ -191,7 +203,11 @@ def __init__( # noqa: C901 resize_param: tuple | None = None target_tmp = target_seg - atlas_seg_ = atlas_seg.apply_pad(((1, 1), (1, 1), (1, 1))) if atlas_seg.is_segmentation_in_border() else atlas_seg + atlas_seg_ = ( + atlas_seg.apply_pad(((1, 1), (1, 1), (1, 1))) + if atlas_seg.is_segmentation_in_border() + else atlas_seg + ) for i in range(_max_iter): if resize_mode == "crop": @@ -200,7 +216,9 @@ def __init__( # noqa: C901 # --- try crop first --- t_crop = target_seg.compute_crop(0, crop_pad_size) - cropped = target_seg.apply_crop(t_crop).apply_pad(crop_pad_size - 50 // 4) + cropped = target_seg.apply_crop(t_crop).apply_pad( + crop_pad_size - 50 // 4 + ) if any(c < o for c, o in zip(cropped.shape, target_seg.shape)): resize_mode = "crop" @@ -224,7 +242,11 @@ def __init__( # noqa: C901 poi_target = poi_target.resample_from_to(target_tmp) if poi_cms is None: - x = atlas_seg_.extract_label(cms_ids, keep_label=True) if cms_ids else atlas_seg_ + x = ( + atlas_seg_.extract_label(cms_ids, keep_label=True) + if cms_ids + else atlas_seg_ + ) poi_cms = calc_centroids(x, second_stage=40, bar=True) if not poi_cms.assert_affine(atlas_seg_, raise_error=False): @@ -242,26 +264,46 @@ def __init__( # noqa: C901 # --- FINAL STEP: apply once to original target --- if resize_mode == "crop": target_seg = target_seg.apply_crop(resize_param) - target_img = target_img.apply_crop(resize_param) if target_img is not None else None + target_img = ( + target_img.apply_crop(resize_param) + if target_img is not None + else None + ) elif resize_mode == "pad": target_seg = target_seg.apply_pad(resize_param) - target_img = target_img.apply_pad(resize_param) if target_img is not None else None - - self.reg_point = Point_Registration(poi_target.resample_from_to(target_seg), poi_cms.resample_from_to(atlas_seg)) + target_img = ( + target_img.apply_pad(resize_param) + if target_img is not None + else None + ) + + self.reg_point = Point_Registration( + poi_target.resample_from_to(target_seg), poi_cms.resample_from_to(atlas_seg) + ) atlas_reg = self.reg_point.transform_nii(atlas_seg, c_val=0) - atlas_img_reg = self.reg_point.transform_nii(atlas_img) if atlas_img is not None else None + atlas_img_reg = ( + self.reg_point.transform_nii(atlas_img) if atlas_img is not None else None + ) if crop: self.crop = (target_seg + atlas_reg).compute_crop(0, 5) target_seg = target_seg.apply_crop(self.crop) - target_img = target_img.apply_crop(self.crop) if target_img is not None else None + target_img = ( + target_img.apply_crop(self.crop) if target_img is not None else None + ) atlas_reg = atlas_reg.apply_crop(self.crop) - atlas_img_reg = atlas_img_reg.apply_crop(self.crop) if atlas_img_reg is not None else None + atlas_img_reg = ( + atlas_img_reg.apply_crop(self.crop) + if atlas_img_reg is not None + else None + ) else: self.crop = None self.target_grid = target_seg.to_gird() - target_seg, atlas_reg, target_img, atlas_img_reg = change_after_point_reg(target_seg, atlas_reg, target_img, atlas_img_reg) + target_seg, atlas_reg, target_img, atlas_img_reg = change_after_point_reg( + target_seg, atlas_reg, target_img, atlas_img_reg + ) self.reg_deform = Deformable_Registration( target_seg if target_img is None else target_img, atlas_reg if atlas_img_reg is None else atlas_img_reg, @@ -350,7 +392,12 @@ def load_(cls, w: tuple) -> Template_Registration: return self - def transform_nii(self, nii_atlas: NII, allow_only_same_grid_as_moving: bool = True, only_rigid=False) -> NII: + def transform_nii( + self, + nii_atlas: NII, + allow_only_same_grid_as_moving: bool = True, + only_rigid=False, + ) -> NII: """Apply both rigid and deformable registration to a NII image. Args: @@ -363,7 +410,9 @@ def transform_nii(self, nii_atlas: NII, allow_only_same_grid_as_moving: bool = T Returns: Transformed ``NII`` aligned with the original target image space. """ - nii_atlas = self.reg_point.transform_nii(nii_atlas, allow_only_same_grid_as_moving=allow_only_same_grid_as_moving) + nii_atlas = self.reg_point.transform_nii( + nii_atlas, allow_only_same_grid_as_moving=allow_only_same_grid_as_moving + ) if only_rigid: return nii_atlas @@ -427,7 +476,9 @@ def transform_poi_inverse(self, poi_target: POI_Global | POI): # poi = poi.apply_crop_inverse(self.crop) # --- inverse rigid point registration --- - poi = self.reg_point.transform_poi_inverse(poi, allow_only_same_grid_as_moving=False) + poi = self.reg_point.transform_poi_inverse( + poi, allow_only_same_grid_as_moving=False + ) # --- back to atlas grid --- poi = poi.resample_from_to(self.atlas_org) @@ -543,18 +594,28 @@ def __init__( # noqa: C901 self.reg_point = pre_registration else: if poi_target_cms is None: - x_seg = target_seg.extract_label(cms_ids, keep_label=True) if cms_ids else target_seg + x_seg = ( + target_seg.extract_label(cms_ids, keep_label=True) + if cms_ids + else target_seg + ) poi_target = calc_centroids(x_seg, second_stage=40, bar=True) else: poi_target = poi_target_cms.resample_from_to(target_seg) if poi_cms is None: - x_seg = atlas_seg.extract_label(cms_ids, keep_label=True) if cms_ids else atlas_seg + x_seg = ( + atlas_seg.extract_label(cms_ids, keep_label=True) + if cms_ids + else atlas_seg + ) poi_cms_local = calc_centroids(x_seg, second_stage=40, bar=True) else: poi_cms_local = poi_cms if not poi_cms_local.assert_affine(atlas_seg, raise_error=False): poi_cms_local = poi_cms_local.resample_from_to(atlas_seg) - self.reg_point = Deepali_Point_Registration(poi_target, poi_cms_local, verbose=False, ddevice=ddevice, gpu=gpu) + self.reg_point = Deepali_Point_Registration( + poi_target, poi_cms_local, verbose=False, ddevice=ddevice, gpu=gpu + ) # --- optional crop ----------------------------------------------------------- # Unlike Template_Registration we do NOT pre-resample the atlas: the rigid @@ -575,11 +636,23 @@ def __init__( # noqa: C901 # on their native grids using same_space=False. We still let the deformable # registration warm-start from the rigid transform by applying the rigid to # the atlas *only for loss evaluation via the deformable pyramid*. - atlas_moved_seg = self.reg_point.transform_nii(atlas_seg, allow_only_same_grid_as_moving=False) - atlas_moved_img = self.reg_point.transform_nii(atlas_img, allow_only_same_grid_as_moving=False) if atlas_img is not None else None + atlas_moved_seg = self.reg_point.transform_nii( + atlas_seg, allow_only_same_grid_as_moving=False + ) + atlas_moved_img = ( + self.reg_point.transform_nii( + atlas_img, allow_only_same_grid_as_moving=False + ) + if atlas_img is not None + else None + ) if crop: atlas_moved_seg = atlas_moved_seg.apply_crop(self.crop) - atlas_moved_img = atlas_moved_img.apply_crop(self.crop) if atlas_moved_img is not None else None + atlas_moved_img = ( + atlas_moved_img.apply_crop(self.crop) + if atlas_moved_img is not None + else None + ) self.reg_deform = Deformable_Registration( target_seg if target_img is None else target_img, @@ -622,12 +695,16 @@ def save(self, path: str | Path) -> None: pickle.dump(self.get_dump(), w) @classmethod - def load(cls, path: str | Path, gpu: int = 0, ddevice: DEVICES = "cuda") -> Template_Registration2: + def load( + cls, path: str | Path, gpu: int = 0, ddevice: DEVICES = "cuda" + ) -> Template_Registration2: with open(path, "rb") as w: return cls.load_(pickle.load(w), gpu=gpu, ddevice=ddevice) @classmethod - def load_(cls, w: tuple, gpu: int = 0, ddevice: DEVICES = "cuda") -> Template_Registration2: + def load_( + cls, w: tuple, gpu: int = 0, ddevice: DEVICES = "cuda" + ) -> Template_Registration2: version, t0, t1, x = w assert version == 1, f"Version mismatch {version=}" self = cls.__new__(cls) @@ -643,9 +720,16 @@ def load_(cls, w: tuple, gpu: int = 0, ddevice: DEVICES = "cuda") -> Template_Re return self # ------------------------------------------------------------------------- warping - def transform_nii(self, nii_atlas: NII, allow_only_same_grid_as_moving: bool = True, only_rigid: bool = False) -> NII: + def transform_nii( + self, + nii_atlas: NII, + allow_only_same_grid_as_moving: bool = True, + only_rigid: bool = False, + ) -> NII: """Warp an atlas NII into the target space (rigid + deformable).""" - nii_atlas = self.reg_point.transform_nii(nii_atlas, allow_only_same_grid_as_moving=allow_only_same_grid_as_moving) + nii_atlas = self.reg_point.transform_nii( + nii_atlas, allow_only_same_grid_as_moving=allow_only_same_grid_as_moving + ) if only_rigid: return nii_atlas if self.crop is not None: @@ -679,6 +763,8 @@ def transform_poi_inverse(self, poi_target: POI_Global | POI) -> POI: poi = poi.resample_from_to(self.target_grid) reg_deform_inv = self.reg_deform.inverse() poi = reg_deform_inv.transform_poi(poi) - poi = self.reg_point.transform_poi_inverse(poi, allow_only_same_grid_as_moving=False) + poi = self.reg_point.transform_poi_inverse( + poi, allow_only_same_grid_as_moving=False + ) poi = poi.resample_from_to(self.atlas_org) return poi diff --git a/TPTBox/segmentation/rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py index 44302c9..3b1c0ca 100644 --- a/TPTBox/segmentation/rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -47,7 +47,9 @@ def _thoracic_like_labels(labels: Iterable[int]) -> list[int]: return [*sorted(v for v in labels if 7 <= int(v) <= 20), 28] -def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> tuple[list[RibCandidate], list[RibCandidate]]: +def _split_ccs_by_side( + rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI +) -> tuple[list[RibCandidate], list[RibCandidate]]: right_axis = cms_vert.get_axis("R") inf_axis = cms_vert.get_axis("I") # ref_x = cms_vert[ref_vert, 50][right_axis] @@ -79,7 +81,9 @@ def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> return left, right -def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None) -> float: +def _touching_surface( + mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None +) -> float: """Return weighted number of voxel faces shared by two masks. Touching along ``up_down_axis`` is weighted by ``up_down_weight``. @@ -101,7 +105,8 @@ def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_wei sl2[axis] = slice(1, None) touching = ( - np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() + np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() + + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() ) weight = axis_weights.get(axis, 1) @@ -137,9 +142,13 @@ def _try_erosion_split( elif _pass == 1: eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) elif _pass == 2: - eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="R") + eroded = cc_nii.erode_msk( + n_pixel=erosion_pixels, verbose=False, ignore_direction="R" + ) elif _pass == 3: - eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="A") + eroded = cc_nii.erode_msk( + n_pixel=erosion_pixels, verbose=False, ignore_direction="A" + ) else: eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) except Exception as e: @@ -158,7 +167,9 @@ def _try_erosion_split( for sub in sub_labels.copy(): count = int((infected_arr == sub).sum()) - should_merge = count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 + should_merge = ( + count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 + ) if not should_merge: continue # print("merge", count, full_size, count / full_size) @@ -173,7 +184,9 @@ def _try_erosion_split( else: target = max( remaining, - key=lambda candidate: _touching_surface(sub_mask, infected_arr == candidate, infected), + key=lambda candidate: _touching_surface( + sub_mask, infected_arr == candidate, infected + ), ) # print("merge of many", target) infected_arr[infected_arr == sub] = target @@ -236,7 +249,9 @@ def split_touching_rib_ccs( next_label = int(rib_cc.max()) + 1 if rib_cc.shape and rib_cc.max() > 0 else 1 u = {int(x) for x in rib_cc.unique() if x != 0} # ribs counte twice (left/right) - expected_number_of_ccs = len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 + expected_number_of_ccs = ( + len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 + ) if expected_number_of_ccs == 0: expected_number_of_ccs = 24 @@ -248,8 +263,12 @@ def split_touching_rib_ccs( binary_ccs = [(cc_label, rib_cc.extract_label(cc_label)) for cc_label in labels] if num_workers == 1: results = [] - for cc_label, binary_cc in tqdm(binary_ccs, total=len(binary_ccs), desc="Separate RIBs"): - result = _try_erosion_split(cc_label, binary_cc, erosion_pixels, min_volume, _pass) + for cc_label, binary_cc in tqdm( + binary_ccs, total=len(binary_ccs), desc="Separate RIBs" + ): + result = _try_erosion_split( + cc_label, binary_cc, erosion_pixels, min_volume, _pass + ) if result is not None: results.append(result) else: @@ -269,7 +288,9 @@ def split_touching_rib_ccs( ] # tqdm advances immediately whenever an individual future completes. - for future in tqdm(as_completed(futures), total=len(futures), desc="Separate RIBs"): + for future in tqdm( + as_completed(futures), total=len(futures), desc="Separate RIBs" + ): result = future.result() if result is not None: results.append(result) @@ -292,7 +313,9 @@ def split_touching_rib_ccs( u.add(new_lbl) if verbose: - logger.print(f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})") + logger.print( + f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})" + ) if short_cut and rib_cc.max() >= expected_number_of_ccs: break @@ -362,12 +385,16 @@ def assign_ribs_to_vert_segmentation( vert_seg = vert_seg.reorient() sem_seg = sem_seg.reorient() rib_seg = rib_seg.reorient(verbose=verbose) - vert_pred = vert_seg.extract_label([Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1]) + vert_pred = vert_seg.extract_label( + [Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1] + ) rib_seg = rib_seg.extract_label([left_id, right_id], keep_label=True) # type: ignore logger.on_debug(f"{rib_seg.unique()=}") # Remove rib voxels overlapping vertebrae rib_seg[vert_pred != 0] = 0 - rib_cc = rib_seg.filter_connected_components(None, min_volume=min_volume, keep_label=False) + rib_cc = rib_seg.filter_connected_components( + None, min_volume=min_volume, keep_label=False + ) cms_vert = calc_centroids(vert_seg) if split_touching: rib_cc = split_touching_rib_ccs( @@ -378,7 +405,9 @@ def assign_ribs_to_vert_segmentation( short_cut=short_cut, verbose=verbose, ) - cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) + cms_cc = calc_centroids( + rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5) + ) cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) vert_labels = _thoracic_like_labels(vert_seg.unique()) @@ -402,12 +431,17 @@ def assign_ribs_to_vert_segmentation( assert isinstance(vert_labels, (tuple, list)), type(vert_labels) if vert_labels and ( - (vert_labels[0] != 28 and vert_labels[1] != 28) or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) + (vert_labels[0] != 28 and vert_labels[1] != 28) + or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) ): inf_axis = cms_vert.get_axis("I") t1_z = cms_vert[vert_labels[0], 50][inf_axis] # use median vertebral spacing as adaptive threshold - spacing = abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) if len(vert_labels) > 1 else 30 + spacing = ( + abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) + if len(vert_labels) > 1 + else 30 + ) threshold = max(5, spacing * 0.6) nearest_left = abs(left_ccs[0].z - t1_z) if left_ccs else float("inf") @@ -434,24 +468,32 @@ def assign_ribs_to_vert_segmentation( logger.on_debug(f"{rib_inst.unique()=}") # Merge rib assignments back into the original vert / sem segmentations # without disturbing existing (non-rib) labels. Skip unmatched CCs (sentinel error_value=255). - print("rib_inst", rib_inst.dtype, rib_inst.seg) matched = np.logical_and(rib_inst != 0, rib_inst != error_value) vert_seg[matched] = rib_inst[matched] - sem_seg[rib_seg != 0] = rib_seg.map_labels({left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value})[rib_seg != 0] # type: ignore + sem_seg[rib_seg != 0] = rib_seg.map_labels( + {left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value} + )[rib_seg != 0] # type: ignore undefined = sum(v == error_value for v in rib_vert_map.values()) logger.print(f"Unmatched rib CCs: {undefined}") if add_error: # or split_touching: vert_seg[rib_inst == error_value] = error_value if np.any(vert_seg.get_seg_array() == error_value): - vert_seg2 = vert_seg.remove_labels(error_value).infect(vert_seg.extract_label(error_value), verbose=False) + vert_seg2 = vert_seg.remove_labels(error_value).infect( + vert_seg.extract_label(error_value), verbose=False + ) vert_seg[vert_seg != vert_seg2] = vert_seg2[vert_seg != vert_seg2] - vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = error_value + vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = ( + error_value + ) vert_seg[ np.logical_and( vert_seg == 0, - sem_seg.extract_label([Location.Rib_Left.value, Location.Rib_Right.value]) == 1, + sem_seg.extract_label( + [Location.Rib_Left.value, Location.Rib_Right.value] + ) + == 1, ) ] = error_value return vert_seg.reorient_(ori), sem_seg.reorient_(ori)