Spenhouet
commented
2 years ago It took me some time but this can be done by
- Getting all corners of the 3D bounding box.
- Projecting all corners into the voxel space of the target image but first projecting into the mm space and then projecting into the target voxel space with the inverse of the target affine.
- Get min. and max. for all corner coordinates to find the new bounding box.
- Round and convert the bounding box values to make them discrete.
- Clip the bounding box by the start and end of the target image.
In addition to the changes proposed in https://github.com/Project-MONAI/MONAI/issues/3167, I made the following changes to make it work:
def get_box(arr: NdarrayOrTensor) -> np.ndarray:
"""For an image get an array of n-dimensional start and end points."""
bottom_corner = np.array(arr.shape)[-3:] # Only use coord dims
return np.array([np.zeros_like(bottom_corner), bottom_corner - 1])
def get_corners(border_box_points: np.ndarray) -> np.ndarray:
return np.array(list(itertools.product(*zip(*border_box_points))))
def project_points(points: np.ndarray, affine_origin: np.ndarray, affine_target: np.ndarray):
points_ = np.insert(points, points.shape[-1], 1, axis=-1)
inv_affine_target = np.linalg.inv(affine_target)
projection = inv_affine_target @ affine_origin
return (points_ @ projection.T)[:, :-1]
def minmax(data, axis=None) -> np.ndarray:
return np.array([np.min(data, axis=axis), np.max(data, axis=axis)])
def project_box(box: np.ndarray,
affine_origin: np.ndarray,
affine_target: np.ndarray,
clip_box: Optional[np.ndarray] = None,
discreate: bool = False) -> np.ndarray:
corners = get_corners(box)
projected_points = project_points(corners, affine_origin, affine_target)
projected_box = minmax(projected_points, axis=0)
if discreate:
projected_box = np.round(projected_box).astype(np.int32)
if clip_box is not None:
projected_box = np.clip(projected_box, a_min=clip_box[0], a_max=clip_box[1])
return projected_box
class CropForegroundd(MapTransform, InvertibleTransform):
"""
Dictionary-based version :py:class:`monai.transforms.CropForeground`.
Crop only the foreground object of the expected images.
The typical usage is to help training and evaluation if the valid part is small in the whole medical image.
The valid part can be determined by any field in the data with `source_key`, for example:
- Select values > 0 in image field as the foreground and crop on all fields specified by `keys`.
- Select label = 3 in label field as the foreground to crop on all fields specified by `keys`.
- Select label > 0 in the third channel of a One-Hot label field as the foreground to crop all `keys` fields.
Users can define arbitrary function to select expected foreground from the whole source image or specified
channels. And it can also add margin to every dim of the bounding box of foreground object.
"""
def __init__(
self,
keys: KeysCollection,
source_key: str,
select_fn: Callable = is_positive,
channel_indices: Optional[IndexSelection] = None,
margin: Union[Sequence[int], int] = 0,
k_divisible: Union[Sequence[int], int] = 1,
mode: NumpyPadModeSequence = NumpyPadMode.CONSTANT,
start_coord_key: str = "foreground_start_coord",
end_coord_key: str = "foreground_end_coord",
allow_missing_keys: bool = False,
**np_kwargs,
) -> None:
"""
Args:
keys: keys of the corresponding items to be transformed.
See also: :py:class:`monai.transforms.compose.MapTransform`
source_key: data source to generate the bounding box of foreground, can be image or label, etc.
select_fn: function to select expected foreground, default is to select values > 0.
channel_indices: if defined, select foreground only on the specified channels
of image. if None, select foreground on the whole image.
margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims.
k_divisible: make each spatial dimension to be divisible by k, default to 1.
if `k_divisible` is an int, the same `k` be applied to all the input spatial dimensions.
mode: padding mode {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``,
``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
one of the listed string values or a user supplied function. Defaults to ``"constant"``.
see also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
it also can be a sequence of string, each element corresponds to a key in ``keys``.
start_coord_key: key to record the start coordinate of spatial bounding box for foreground.
end_coord_key: key to record the end coordinate of spatial bounding box for foreground.
allow_missing_keys: don't raise exception if key is missing.
np_kwargs: other args for `np.pad` API, note that `np.pad` treats channel dimension as the first dimension.
more details: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
"""
super().__init__(keys, allow_missing_keys)
self.source_key = source_key
self.start_coord_key = start_coord_key
self.end_coord_key = end_coord_key
self.cropper = CropForeground(
select_fn=select_fn,
channel_indices=channel_indices,
margin=margin,
k_divisible=k_divisible,
**np_kwargs,
)
self.mode = ensure_tuple_rep(mode, len(self.keys))
def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Dict[Hashable, np.ndarray]:
d = dict(data)
img: np.ndarray
img, *_ = convert_data_type(d[self.source_key], np.ndarray) # type: ignore
source_meta_key = f"{self.source_key}_meta_dict"
bounding_box = self.cropper.compute_bounding_box(img)
source_affine = d[source_meta_key]['affine']
d[self.start_coord_key] = bounding_box[0]
d[self.end_coord_key] = bounding_box[1]
for key, m in self.key_iterator(d, self.mode):
meta_key = f"{key}_meta_dict"
box = project_box(
np.array(bounding_box),
affine_origin=source_affine,
affine_target=d[meta_key]['affine'],
clip_box=get_box(d[key]),
discreate=True,
)
self.push_transform(d, key, extra_info={"box_start": box[0], "box_end": box[1]})
d[key], d[meta_key] = self.cropper(img=d[key],
box_start=box[0],
box_end=box[1],
mode=m,
meta_data=d[meta_key])
return d
Now the results are correct.
Nic-Ma
ericspod
rijobro
vikashg
Describe the bug The bounding box for the crop is not projected into the space of the target image. This leads to wrong crops if the target image has another FoV or orientation than the image where the bounding box was taken.
To Reproduce
import numpy as np from monai.transforms.compose import Compose
Use the fixed transform as proposed in #3167
from path.to.the.fixed.transform import CropForegroundd from monai.transforms.io.dictionary import LoadImaged, SaveImaged from monai.transforms.utility.dictionary import EnsureChannelFirstd, ToTensord
path = Path('scratch/monai_crop_bug')
SEG_KEY = 'seg' THALA_KEY = 'thalamic' FILE_KEYS = [SEG_KEY, THALA_KEY]
data = { SEG_KEY: path / 'seg.nii.gz', THALA_KEY: path / 'thalamic.nii.gz', }
margin = 20
process = Compose([ LoadImaged(FILE_KEYS), ToTensord(FILE_KEYS), EnsureChannelFirstd(FILE_KEYS), CropForegroundd(FILE_KEYS, source_key=SEG_KEY, margin=margin), SaveImaged( FILE_KEYS, output_dir='scratch/monai_crop_bug', outputpostfix=f"crop{margin}", resample=False, output_dtype=np.int8, separate_folder=False, ), ])
results = process([data])