|
|
import numpy as np |
|
|
|
|
|
|
|
|
|
|
|
from acvl_utils.cropping_and_padding.bounding_boxes import get_bbox_from_mask, crop_to_bbox, bounding_box_to_slice |
|
|
|
|
|
|
|
|
def create_nonzero_mask(data): |
|
|
""" |
|
|
|
|
|
:param data: |
|
|
:return: the mask is True where the data is nonzero |
|
|
""" |
|
|
from scipy.ndimage import binary_fill_holes |
|
|
assert data.ndim in (3, 4), "data must have shape (C, X, Y, Z) or shape (C, X, Y)" |
|
|
nonzero_mask = np.zeros(data.shape[1:], dtype=bool) |
|
|
for c in range(data.shape[0]): |
|
|
this_mask = data[c] != 0 |
|
|
nonzero_mask = nonzero_mask | this_mask |
|
|
nonzero_mask = binary_fill_holes(nonzero_mask) |
|
|
return nonzero_mask |
|
|
|
|
|
|
|
|
def crop_to_nonzero(data, seg=None, nonzero_label=-1): |
|
|
""" |
|
|
|
|
|
:param data: |
|
|
:param seg: |
|
|
:param nonzero_label: this will be written into the segmentation map |
|
|
:return: |
|
|
""" |
|
|
nonzero_mask = create_nonzero_mask(data) |
|
|
nonzero_mask = np.ones(nonzero_mask.shape, dtype=bool) |
|
|
|
|
|
bbox = get_bbox_from_mask(nonzero_mask) |
|
|
|
|
|
slicer = bounding_box_to_slice(bbox) |
|
|
data = data[tuple([slice(None), *slicer])] |
|
|
|
|
|
if seg is not None: |
|
|
seg = seg[tuple([slice(None), *slicer])] |
|
|
|
|
|
nonzero_mask = nonzero_mask[slicer][None] |
|
|
if seg is not None: |
|
|
seg[(seg == 0) & (~nonzero_mask)] = nonzero_label |
|
|
else: |
|
|
nonzero_mask = nonzero_mask.astype(np.int8) |
|
|
nonzero_mask[nonzero_mask == 0] = nonzero_label |
|
|
nonzero_mask[nonzero_mask > 0] = 0 |
|
|
seg = nonzero_mask |
|
|
return data, seg, bbox |
|
|
|
|
|
|
|
|
|