Skip to content

bug fix _rib_assingn #158

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
TanjaLe wants to merge 3 commits into from
Closed

bug fix _rib_assingn #158

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
449ebcd
bug fix _rib_assingn
TanjaLe Sep 30, 2026
41e4a15
style fixes by ruff
TanjaLe Sep 30, 2026
2716437
deepali point registration and clean up
TanjaLe Oct 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
146 changes: 116 additions & 30 deletions TPTBox/registration/_deformable/multilabel_segmentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -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":
Expand All @@ -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"
Expand All @@ -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):
Expand All @@ -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,
Expand Down Expand Up @@ -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:
Expand All @@ -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

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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
Loading
Loading