File size: 1,616 Bytes
19c1f58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import numpy as np
# Hello! crop_to_nonzero is the function you are looking for. Ignore the rest.
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) #arthur : disable cropping in preprocessing
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
|