# Copyright (c) MONAI Consortium

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

from __future__ import annotations

import json

from collections.abc import Callable, Hashable, Iterable, Sequence

from typing import Any

import numpy as np

import torch

from monai.config import IndexSelection, KeysCollection, NdarrayOrTensor

from monai.networks.layers import GaussianFilter

from monai.transforms import Resize, SpatialCrop

from monai.transforms.transform import MapTransform, Randomizable, Transform

from monai.transforms.utils import generate_spatial_bounding_box, is_positive

from monai.utils import InterpolateMode, ensure_tuple, ensure_tuple_rep, min_version, optional_import

from monai.utils.enums import PostFix

measure, _ = optional_import("skimage.measure", "0.14.2", min_version)

distance_transform_cdt, _ = optional_import("scipy.ndimage", name="distance_transform_cdt")

DEFAULT_POST_FIX = PostFix.meta()

# Transforms to support Training for Deepgrow models

class FindAllValidSlicesd(Transform):

"""

Find/List all valid slices in the label.

Label is assumed to be a 4D Volume with shape CDHW, where C=1.

Args:

label: key to the label source.

sids: key to store slices indices having valid label map.

"""

def __init__(self, label: str = "label", sids: str = "sids"):

self.label = label

self.sids = sids

def _apply(self, label):

sids = []

for sid in range(label.shape[1]): # Assume channel is first

if np.sum(label[0][sid]) != 0:

sids.append(sid)

return np.asarray(sids)

def __call__(self, data: Any) -> dict:

d: dict = dict(data)

label = d[self.label].numpy() if isinstance(data[self.label], torch.Tensor) else data[self.label]

if label.shape[0] != 1:

raise ValueError(f"Only supports single channel labels, got label shape {label.shape}!")

if len(label.shape) != 4: # only for 3D

raise ValueError(f"Only supports label with shape CDHW, got label shape {label.shape}!")

sids = self._apply(label)

if sids is not None and len(sids):

d[self.sids] = sids

return d

class AddInitialSeedPointd(Randomizable, Transform):

"""

Add random guidance as initial seed point for a given label.

Note that the label is of size (C, D, H, W) or (C, H, W)

The guidance is of size (2, N, # of dims) where N is number of guidance added.

# of dims = 4 when C, D, H, W; # of dims = 3 when (C, H, W)

Args:

label: label source.

guidance: key to store guidance.

sids: key that represents list of valid slice indices for the given label.

sid: key that represents the slice to add initial seed point. If not present, random sid will be chosen.

connected_regions: maximum connected regions to use for adding initial points.

"""

def __init__(

self,

label: str = "label",

guidance: str = "guidance",

sids: str = "sids",

sid: str = "sid",

connected_regions: int = 5,

):

self.label = label

self.sids_key = sids

self.sid_key = sid

self.sid = None

self.guidance = guidance

self.connected_regions = connected_regions

def randomize(self, data):

sid = data.get(self.sid_key, None)

sids = data.get(self.sids_key, None)

if sids is not None:

if sid is None or sid not in sids:

sid = self.R.choice(sids, replace=False)

else:

sid = None

self.sid = sid

def _apply(self, label, sid):

dimensions = 3 if len(label.shape) > 3 else 2

default_guidance = [-1] * (dimensions + 1)

dims = dimensions

if sid is not None and dimensions == 3:

dims = 2

label = label[0][sid][np.newaxis] # Assume channel is first

label = (label > 0.5).astype(np.float32)

blobs_labels = measure.label(label.astype(int), background=0) if dims == 2 else label

if np.max(blobs_labels) <= 0:

raise AssertionError("Not a valid Label")

pos_guidance = []

for ridx in range(1, 2 if dims == 3 else self.connected_regions + 1):

if dims == 2:

label = (blobs_labels == ridx).astype(np.float32)

if np.sum(label) == 0:

pos_guidance.append(default_guidance)

continue

distance = distance_transform_cdt(label).flatten()

probability = np.exp(distance) - 1.0

idx = np.where(label.flatten() > 0)[0]

seed = self.R.choice(idx, size=1, p=probability[idx] / np.sum(probability[idx]))

dst = distance[seed]

g = np.asarray(np.unravel_index(seed, label.shape)).transpose().tolist()[0]

g[0] = dst[0] # for debug

if dimensions == 2 or dims == 3:

pos_guidance.append(g)

else:

pos_guidance.append([g[0], sid, g[-2], g[-1]])

return np.asarray([pos_guidance, [default_guidance] * len(pos_guidance)])

def __call__(self, data):

d = dict(data)

self.randomize(data)

d[self.guidance] = json.dumps(self._apply(d[self.label], self.sid).astype(int, copy=False).tolist())

return d

class AddGuidanceSignald(Transform):

"""

Add Guidance signal for input image.

Based on the "guidance" points, apply gaussian to them and add them as new channel for input image.

Args:

image: key to the image source.

guidance: key to store guidance.

sigma: standard deviation for Gaussian kernel.

number_intensity_ch: channel index.

"""

def __init__(self, image: str = "image", guidance: str = "guidance", sigma: int = 2, number_intensity_ch: int = 1):

self.image = image

self.guidance = guidance

self.sigma = sigma

self.number_intensity_ch = number_intensity_ch

def _get_signal(self, image, guidance):

dimensions = 3 if len(image.shape) > 3 else 2

guidance = guidance.tolist() if isinstance(guidance, np.ndarray) else guidance

guidance = json.loads(guidance) if isinstance(guidance, str) else guidance

if dimensions == 3:

signal = np.zeros((len(guidance), image.shape[-3], image.shape[-2], image.shape[-1]), dtype=np.float32)

else:

signal = np.zeros((len(guidance), image.shape[-2], image.shape[-1]), dtype=np.float32)

sshape = signal.shape

for i, g_i in enumerate(guidance):

for point in g_i:

if np.any(np.asarray(point) < 0):

continue

if dimensions == 3:

p1 = max(0, min(int(point[-3]), sshape[-3] - 1))

p2 = max(0, min(int(point[-2]), sshape[-2] - 1))

p3 = max(0, min(int(point[-1]), sshape[-1] - 1))

signal[i, p1, p2, p3] = 1.0

else:

p1 = max(0, min(int(point[-2]), sshape[-2] - 1))

p2 = max(0, min(int(point[-1]), sshape[-1] - 1))

signal[i, p1, p2] = 1.0

if np.max(signal[i]) > 0:

signal_tensor = torch.tensor(signal[i])

pt_gaussian = GaussianFilter(len(signal_tensor.shape), sigma=self.sigma)

signal_tensor = pt_gaussian(signal_tensor.unsqueeze(0).unsqueeze(0))

signal_tensor = signal_tensor.squeeze(0).squeeze(0)

signal[i] = signal_tensor.detach().cpu().numpy()

signal[i] = (signal[i] - np.min(signal[i])) / (np.max(signal[i]) - np.min(signal[i]))

return signal

def _apply(self, image, guidance):

signal = self._get_signal(image, guidance)

if isinstance(image, torch.Tensor):

image = image.detach().cpu().numpy()

image = image[0 : 0 + self.number_intensity_ch, ...]

return np.concatenate([image, signal], axis=0)

def __call__(self, data):

d = dict(data)

image = d[self.image]

guidance = d[self.guidance]

d[self.image] = self._apply(image, guidance)

return d

class FindDiscrepancyRegionsd(Transform):

"""

Find discrepancy between prediction and actual during click interactions during training.

Args:

label: key to label source.

pred: key to prediction source.

discrepancy: key to store discrepancies found between label and prediction.

"""

def __init__(self, label: str = "label", pred: str = "pred", discrepancy: str = "discrepancy"):

self.label = label

self.pred = pred

self.discrepancy = discrepancy

@staticmethod

def disparity(label, pred):

label = (label > 0.5).astype(np.float32)

pred = (pred > 0.5).astype(np.float32)

disparity = label - pred

pos_disparity = (disparity > 0).astype(np.float32)

neg_disparity = (disparity < 0).astype(np.float32)

return [pos_disparity, neg_disparity]

def _apply(self, label, pred):

return self.disparity(label, pred)

def __call__(self, data):

d = dict(data)

label = d[self.label]

pred = d[self.pred]

d[self.discrepancy] = self._apply(label, pred)

return d

class AddRandomGuidanced(Randomizable, Transform):

"""

Add random guidance based on discrepancies that were found between label and prediction.

input shape is as below:

Guidance is of shape (2, N, # of dim)

Discrepancy is of shape (2, C, D, H, W) or (2, C, H, W)

Probability is of shape (1)

Args:

guidance: key to guidance source.

discrepancy: key that represents discrepancies found between label and prediction.

probability: key that represents click/interaction probability.

"""

def __init__(self, guidance: str = "guidance", discrepancy: str = "discrepancy", probability: str = "probability"):

self.guidance = guidance

self.discrepancy = discrepancy

self.probability = probability

self._will_interact = None

def randomize(self, data=None):

probability = data[self.probability]

self._will_interact = self.R.choice([True, False], p=[probability, 1.0 - probability])

def find_guidance(self, discrepancy):

distance = distance_transform_cdt(discrepancy).flatten()

probability = np.exp(distance) - 1.0

idx = np.where(discrepancy.flatten() > 0)[0]

if np.sum(discrepancy > 0) > 0:

seed = self.R.choice(idx, size=1, p=probability[idx] / np.sum(probability[idx]))

dst = distance[seed]

g = np.asarray(np.unravel_index(seed, discrepancy.shape)).transpose().tolist()[0]

g[0] = dst[0]

return g

return None

def add_guidance(self, discrepancy, will_interact):

if not will_interact:

return None, None

pos_discr = discrepancy[0]

neg_discr = discrepancy[1]

can_be_positive = np.sum(pos_discr) > 0

can_be_negative = np.sum(neg_discr) > 0

correct_pos = np.sum(pos_discr) >= np.sum(neg_discr)

if correct_pos and can_be_positive:

return self.find_guidance(pos_discr), None

if not correct_pos and can_be_negative:

return None, self.find_guidance(neg_discr)

return None, None

def _apply(self, guidance, discrepancy):

guidance = guidance.tolist() if isinstance(guidance, np.ndarray) else guidance

guidance = json.loads(guidance) if isinstance(guidance, str) else guidance

pos, neg = self.add_guidance(discrepancy, self._will_interact)

if pos:

guidance[0].append(pos)

guidance[1].append([-1] * len(pos))

if neg:

guidance[0].append([-1] * len(neg))

guidance[1].append(neg)

return json.dumps(np.asarray(guidance, dtype=int).tolist())

def __call__(self, data):

d = dict(data)

guidance = d[self.guidance]

discrepancy = d[self.discrepancy]

self.randomize(data)

d[self.guidance] = self._apply(guidance, discrepancy)

return d

class SpatialCropForegroundd(MapTransform):

"""

Crop only the foreground object of the expected images.

Difference VS :py:class:`monai.transforms.CropForegroundd`:

1. If the bounding box is smaller than spatial size in all dimensions then this transform will crop the

object using box's center and spatial_size.

2. This transform will set "start_coord_key", "end_coord_key", "original_shape_key" and "cropped_shape_key"

in data[{key}_{meta_key_postfix}]

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.

Args:

keys: keys of the corresponding items to be transformed.

See also: :py:class:`monai.transforms.MapTransform`

source_key: data source to generate the bounding box of foreground, can be image or label, etc.

spatial_size: minimal spatial size of the image patch e.g. [128, 128, 128] to fit in.

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.

allow_smaller: when computing box size with `margin`, whether allow the image size to be smaller

than box size, default to `True`. if the margined size is bigger than image size, will pad with

specified `mode`.

meta_keys: explicitly indicate the key of the corresponding metadata dictionary.

for example, for data with key `image`, the metadata by default is in `image_meta_dict`.

the metadata is a dictionary object which contains: filename, original_shape, etc.

it can be a sequence of string, map to the `keys`.

if None, will try to construct meta_keys by `key_{meta_key_postfix}`.

meta_key_postfix: if meta_keys is None, use `{key}_{meta_key_postfix}` to fetch/store the metadata according

to the key data, default is `meta_dict`, the metadata is a dictionary object.

For example, to handle key `image`, read/write affine matrices from the

metadata `image_meta_dict` dictionary's `affine` field.

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.

original_shape_key: key to record original shape for foreground.

cropped_shape_key: key to record cropped shape for foreground.

allow_missing_keys: don't raise exception if key is missing.

"""

def __init__(

self,

keys: KeysCollection,

source_key: str,

spatial_size: Sequence[int] | np.ndarray,

select_fn: Callable = is_positive,

channel_indices: IndexSelection | None = None,

margin: int = 0,

allow_smaller: bool = True,

meta_keys: KeysCollection | None = None,

meta_key_postfix: str = DEFAULT_POST_FIX,

start_coord_key: str = "foreground_start_coord",

end_coord_key: str = "foreground_end_coord",

original_shape_key: str = "foreground_original_shape",

cropped_shape_key: str = "foreground_cropped_shape",

allow_missing_keys: bool = False,

) -> None:

super().__init__(keys, allow_missing_keys)

self.source_key = source_key

self.spatial_size = list(spatial_size)

self.select_fn = select_fn

self.channel_indices = channel_indices

self.margin = margin

self.allow_smaller = allow_smaller

self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys)

if len(self.keys) != len(self.meta_keys):

raise ValueError("meta_keys should have the same length as keys.")

self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))

self.start_coord_key = start_coord_key

self.end_coord_key = end_coord_key

self.original_shape_key = original_shape_key

self.cropped_shape_key = cropped_shape_key

def __call__(self, data):

d = dict(data)

box_start, box_end = generate_spatial_bounding_box(

d[self.source_key], self.select_fn, self.channel_indices, self.margin, self.allow_smaller

)

center = list(np.mean([box_start, box_end], axis=0).astype(int, copy=False))

current_size = list(np.subtract(box_end, box_start).astype(int, copy=False))

if np.all(np.less(current_size, self.spatial_size)):

cropper = SpatialCrop(roi_center=center, roi_size=self.spatial_size)

box_start = [s.start for s in cropper.slices]

box_end = [s.stop for s in cropper.slices]

else:

cropper = SpatialCrop(roi_start=box_start, roi_end=box_end)

for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix):

meta_key = meta_key or f"{key}_{meta_key_postfix}"

d[meta_key][self.start_coord_key] = box_start

d[meta_key][self.end_coord_key] = box_end

d[meta_key][self.original_shape_key] = d[key].shape

image = cropper(d[key])

d[meta_key][self.cropped_shape_key] = image.shape

d[key] = image

return d

# Transforms to support Inference for Deepgrow models

class AddGuidanceFromPointsd(Transform):

"""

Add guidance based on user clicks.

We assume the input is loaded by LoadImaged and has the shape of (H, W, D) originally.

Clicks always specify the coordinates in (H, W, D)

If depth_first is True:

Input is now of shape (D, H, W), will return guidance that specifies the coordinates in (D, H, W)

else:

Input is now of shape (H, W, D), will return guidance that specifies the coordinates in (H, W, D)

Args:

ref_image: key to reference image to fetch current and original image details.

guidance: output key to store guidance.

foreground: key that represents user foreground (+ve) clicks.

background: key that represents user background (-ve) clicks.

axis: axis that represents slices in 3D volume. (axis to Depth)

depth_first: if depth (slices) is positioned at first dimension.

spatial_dims: dimensions based on model used for deepgrow (2D vs 3D).

slice_key: key that represents applicable slice to add guidance.

meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`.

for example, for data with key `image`, the metadata by default is in `image_meta_dict`.

the metadata is a dictionary object which contains: filename, original_shape, etc.

if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`.

meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the metadata according

to the key data, default is `meta_dict`, the metadata is a dictionary object.

For example, to handle key `image`, read/write affine matrices from the

metadata `image_meta_dict` dictionary's `affine` field.

"""

def __init__(

self,

ref_image: str,

guidance: str = "guidance",

foreground: str = "foreground",

background: str = "background",

axis: int = 0,

depth_first: bool = True,

spatial_dims: int = 2,

slice_key: str = "slice",

meta_keys: str | None = None,

meta_key_postfix: str = DEFAULT_POST_FIX,

):

self.ref_image = ref_image

self.guidance = guidance

self.foreground = foreground

self.background = background

self.axis = axis

self.depth_first = depth_first

self.dimensions = spatial_dims

self.slice = slice_key

self.meta_keys = meta_keys

self.meta_key_postfix = meta_key_postfix

def _apply(self, pos_clicks, neg_clicks, factor, slice_num):

pos = neg = []

if self.dimensions == 2:

points: list = list(pos_clicks)

points.extend(neg_clicks)

slices = list(np.unique(np.array(points)[:, self.axis]))

slice_idx = slices[0] if slice_num is None else next(x for x in slices if x == slice_num)

if len(pos_clicks):

pos_clicks = np.array(pos_clicks)

pos = (pos_clicks[np.where(pos_clicks[:, self.axis] == slice_idx)] * factor)[:, 1:].astype(int).tolist()

if len(neg_clicks):

neg_clicks = np.array(neg_clicks)

neg = (neg_clicks[np.where(neg_clicks[:, self.axis] == slice_idx)] * factor)[:, 1:].astype(int).tolist()

guidance = [pos, neg, slice_idx]

else:

if len(pos_clicks):

pos = np.multiply(pos_clicks, factor).astype(int, copy=False).tolist()

if len(neg_clicks):

neg = np.multiply(neg_clicks, factor).astype(int, copy=False).tolist()

guidance = [pos, neg]

return guidance

def __call__(self, data):

d = dict(data)

meta_dict_key = self.meta_keys or f"{self.ref_image}_{self.meta_key_postfix}"

if meta_dict_key not in d:

raise RuntimeError(f"Missing meta_dict {meta_dict_key} in data!")

if "spatial_shape" not in d[meta_dict_key]:

raise RuntimeError('Missing "spatial_shape" in meta_dict!')

original_shape = d[meta_dict_key]["spatial_shape"]

current_shape = list(d[self.ref_image].shape)

if self.depth_first:

if self.axis != 0:

raise RuntimeError("Depth first means the depth axis should be 0.")

# in here we assume the depth dimension was in the last dimension of "original_shape"

original_shape = np.roll(original_shape, 1)

factor = np.array(current_shape) / original_shape

fg_bg_clicks = []

for key in [self.foreground, self.background]:

clicks = d[key]

clicks = list(np.array(clicks, dtype=int))

if self.depth_first:

for i in range(len(clicks)):

clicks[i] = list(np.roll(clicks[i], 1))

fg_bg_clicks.append(clicks)

d[self.guidance] = self._apply(fg_bg_clicks[0], fg_bg_clicks[1], factor, d.get(self.slice))

return d

class SpatialCropGuidanced(MapTransform):

"""

Crop image based on guidance with minimal spatial size.

- If the bounding box is smaller than spatial size in all dimensions then this transform will crop the

object using box's center and spatial_size.

- This transform will set "start_coord_key", "end_coord_key", "original_shape_key" and "cropped_shape_key"

in data[{key}_{meta_key_postfix}]

Input data is of shape (C, spatial_1, [spatial_2, ...])

Args:

keys: keys of the corresponding items to be transformed.

guidance: key to the guidance. It is used to generate the bounding box of foreground

spatial_size: minimal spatial size of the image patch e.g. [128, 128, 128] to fit in.

margin: add margin value to spatial dims of the bounding box, if only 1 value provided, use it for all dims.

meta_keys: explicitly indicate the key of the corresponding metadata dictionary.

for example, for data with key `image`, the metadata by default is in `image_meta_dict`.

the metadata is a dictionary object which contains: filename, original_shape, etc.

it can be a sequence of string, map to the `keys`.

if None, will try to construct meta_keys by `key_{meta_key_postfix}`.

meta_key_postfix: if meta_keys is None, use `key_{postfix}` to fetch the metadata according

to the key data, default is `meta_dict`, the metadata is a dictionary object.

For example, to handle key `image`, read/write affine matrices from the

metadata `image_meta_dict` dictionary's `affine` field.

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.

original_shape_key: key to record original shape for foreground.

cropped_shape_key: key to record cropped shape for foreground.

allow_missing_keys: don't raise exception if key is missing.

"""

def __init__(

self,

keys: KeysCollection,

guidance: str,

spatial_size: Iterable[int],

margin: int = 20,

meta_keys: KeysCollection | None = None,

meta_key_postfix: str = DEFAULT_POST_FIX,

start_coord_key: str = "foreground_start_coord",

end_coord_key: str = "foreground_end_coord",

original_shape_key: str = "foreground_original_shape",

cropped_shape_key: str = "foreground_cropped_shape",

allow_missing_keys: bool = False,

) -> None:

super().__init__(keys, allow_missing_keys)

self.guidance = guidance

self.spatial_size = list(spatial_size)

self.margin = margin

self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys)

if len(self.keys) != len(self.meta_keys):

raise ValueError("meta_keys should have the same length as keys.")

self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))

self.start_coord_key = start_coord_key

self.end_coord_key = end_coord_key

self.original_shape_key = original_shape_key

self.cropped_shape_key = cropped_shape_key

def bounding_box(self, points, img_shape):

ndim = len(img_shape)

margin = ensure_tuple_rep(self.margin, ndim)

for m in margin:

if m < 0:

raise ValueError("margin value should not be negative number.")

box_start = [0] * ndim

box_end = [0] * ndim

for di in range(ndim):

dt = points[..., di]

min_d = max(min(dt - margin[di]), 0)

max_d = min(img_shape[di], max(dt + margin[di] + 1))

box_start[di], box_end[di] = min_d, max_d

return box_start, box_end

def __call__(self, data: Any) -> dict:

d: dict = dict(data)

first_key: Hashable = self.first_key(d)

if first_key == ():

return d

guidance = d[self.guidance]

original_spatial_shape = d[first_key].shape[1:]

box_start, box_end = self.bounding_box(np.array(guidance[0] + guidance[1]), original_spatial_shape)

center = list(np.mean([box_start, box_end], axis=0).astype(int, copy=False))

spatial_size = self.spatial_size

box_size = list(np.subtract(box_end, box_start).astype(int, copy=False))

spatial_size = spatial_size[-len(box_size) :]

if len(spatial_size) < len(box_size):

# If the data is in 3D and spatial_size is specified as 2D [256,256]

# Then we will get all slices in such case

diff = len(box_size) - len(spatial_size)

spatial_size = list(original_spatial_shape[1 : (1 + diff)]) + spatial_size

if np.all(np.less(box_size, spatial_size)):

if len(center) == 3:

# 3D Deepgrow: set center to be middle of the depth dimension (D)

center[0] = spatial_size[0] // 2

cropper = SpatialCrop(roi_center=center, roi_size=spatial_size)

else:

cropper = SpatialCrop(roi_start=box_start, roi_end=box_end)

# update bounding box in case it was corrected by the SpatialCrop constructor

box_start = np.array([s.start for s in cropper.slices])

box_end = np.array([s.stop for s in cropper.slices])

for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix):

if not np.array_equal(d[key].shape[1:], original_spatial_shape):

raise RuntimeError("All the image specified in keys should have same spatial shape")

meta_key = meta_key or f"{key}_{meta_key_postfix}"

d[meta_key][self.start_coord_key] = box_start

d[meta_key][self.end_coord_key] = box_end

d[meta_key][self.original_shape_key] = d[key].shape

image = cropper(d[key])

d[meta_key][self.cropped_shape_key] = image.shape

d[key] = image

pos_clicks, neg_clicks = guidance[0], guidance[1]

pos = np.subtract(pos_clicks, box_start).tolist() if len(pos_clicks) else []

neg = np.subtract(neg_clicks, box_start).tolist() if len(neg_clicks) else []

d[self.guidance] = [pos, neg]

return d

class ResizeGuidanced(Transform):

"""

Resize the guidance based on cropped vs resized image.

This transform assumes that the images have been cropped and resized. And the shape after cropped is store inside

the meta dict of ref image.

Args:

guidance: key to guidance

ref_image: key to reference image to fetch current and original image details

meta_keys: explicitly indicate the key of the metadata dictionary of `ref_image`.

for example, for data with key `image`, the metadata by default is in `image_meta_dict`.

the metadata is a dictionary object which contains: filename, original_shape, etc.

if None, will try to construct meta_keys by `{ref_image}_{meta_key_postfix}`.

meta_key_postfix: if meta_key is None, use `{ref_image}_{meta_key_postfix}` to fetch the metadata according

to the key data, default is `meta_dict`, the metadata is a dictionary object.

For example, to handle key `image`, read/write affine matrices from the

metadata `image_meta_dict` dictionary's `affine` field.

cropped_shape_key: key that records cropped shape for foreground.

"""

def __init__(

self,

guidance: str,

ref_image: str,

meta_keys: str | None = None,

meta_key_postfix: str = DEFAULT_POST_FIX,

cropped_shape_key: str = "foreground_cropped_shape",

) -> None:

self.guidance = guidance

self.ref_image = ref_image

self.meta_keys = meta_keys

self.meta_key_postfix = meta_key_postfix

self.cropped_shape_key = cropped_shape_key

def __call__(self, data: Any) -> dict:

d = dict(data)

guidance = d[self.guidance]

meta_dict: dict = d[self.meta_keys or f"{self.ref_image}_{self.meta_key_postfix}"]

current_shape = d[self.ref_image].shape[1:]

cropped_shape = meta_dict[self.cropped_shape_key][1:]

factor = np.divide(current_shape, cropped_shape)

pos_clicks, neg_clicks = guidance[0], guidance[1]

pos = np.multiply(pos_clicks, factor).astype(int, copy=False).tolist() if len(pos_clicks) else []

neg = np.multiply(neg_clicks, factor).astype(int, copy=False).tolist() if len(neg_clicks) else []

d[self.guidance] = [pos, neg]

return d

class RestoreLabeld(MapTransform):

"""

Restores label based on the ref image.

The ref_image is assumed that it went through the following transforms:

1. Fetch2DSliced (If 2D)

2. Spacingd

3. SpatialCropGuidanced

4. Resized

And its shape is assumed to be (C, D, H, W)

This transform tries to undo these operation so that the result label can be overlapped with original volume.

It does the following operation:

1. Undo Resized

2. Undo SpatialCropGuidanced

3. Undo Spacingd

4. Undo Fetch2DSliced

The resulting label is of shape (D, H, W)

Args:

keys: keys of the corresponding items to be transformed.

ref_image: reference image to fetch current and original image details

slice_only: apply only to an applicable slice, in case of 2D model/prediction

mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``,

``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}

One of the listed string values or a user supplied function for padding. Defaults to ``"constant"``.

See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html

align_corners: Geometrically, we consider the pixels of the input as squares rather than points.

See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html

It also can be a sequence of bool, each element corresponds to a key in ``keys``.

meta_keys: explicitly indicate the key of the corresponding metadata dictionary.

for example, for data with key `image`, the metadata by default is in `image_meta_dict`.

the metadata is a dictionary object which contains: filename, original_shape, etc.

it can be a sequence of string, map to the `keys`.

if None, will try to construct meta_keys by `key_{meta_key_postfix}`.

meta_key_postfix: if meta_key is None, use `key_{meta_key_postfix} to fetch the metadata according

to the key data, default is `meta_dict`, the metadata is a dictionary object.

For example, to handle key `image`, read/write affine matrices from the

metadata `image_meta_dict` dictionary's `affine` field.

start_coord_key: key that records the start coordinate of spatial bounding box for foreground.

end_coord_key: key that records the end coordinate of spatial bounding box for foreground.

original_shape_key: key that records original shape for foreground.

cropped_shape_key: key that records cropped shape for foreground.

allow_missing_keys: don't raise exception if key is missing.

restore_resizing: used to enable or disable resizing restoration, default is True.

If True, the transform will resize the items back to its original shape.

restore_cropping: used to enable or disable cropping restoration, default is True.

If True, the transform will restore the items to its uncropped size.

restore_spacing: used to enable or disable spacing restoration, default is True.

If True, the transform will resample the items back to the spacing it had before being altered.

restore_slicing: used to enable or disable slicing restoration, default is True.

If True, the transform will reassemble the full volume by restoring the slices to their original positions.

"""

def __init__(

self,

keys: KeysCollection,

ref_image: str,

slice_only: bool = False,

mode: Sequence[InterpolateMode | str] | InterpolateMode | str = InterpolateMode.NEAREST,

align_corners: Sequence[bool | None] | bool | None = None,

meta_keys: str | None = None,

meta_key_postfix: str = DEFAULT_POST_FIX,

start_coord_key: str = "foreground_start_coord",

end_coord_key: str = "foreground_end_coord",

original_shape_key: str = "foreground_original_shape",

cropped_shape_key: str = "foreground_cropped_shape",

allow_missing_keys: bool = False,

restore_resizing: bool = True,

restore_cropping: bool = True,

restore_spacing: bool = True,

restore_slicing: bool = True,

) -> None:

super().__init__(keys, allow_missing_keys)

self.ref_image = ref_image

self.slice_only = slice_only

self.mode = ensure_tuple_rep(mode, len(self.keys))

self.align_corners = ensure_tuple_rep(align_corners, len(self.keys))

self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys)

if len(self.keys) != len(self.meta_keys):

raise ValueError("meta_keys should have the same length as keys.")

self.meta_key_postfix = meta_key_postfix

self.start_coord_key = start_coord_key

self.end_coord_key = end_coord_key

self.original_shape_key = original_shape_key

self.cropped_shape_key = cropped_shape_key

self.restore_resizing = restore_resizing

self.restore_cropping = restore_cropping

self.restore_spacing = restore_spacing

self.restore_slicing = restore_slicing

def __call__(self, data: Any) -> dict:

d = dict(data)

meta_dict: dict = d[f"{self.ref_image}_{self.meta_key_postfix}"]

for key, mode, align_corners, meta_key in self.key_iterator(d, self.mode, self.align_corners, self.meta_keys):

image = d[key]

# Undo Resize

if self.restore_resizing:

current_shape = image.shape

cropped_shape = meta_dict[self.cropped_shape_key]

if np.any(np.not_equal(current_shape, cropped_shape)):

resizer = Resize(spatial_size=cropped_shape[1:], mode=mode)

image = resizer(image, mode=mode, align_corners=align_corners)

# Undo Crop

if self.restore_cropping:

original_shape = meta_dict[self.original_shape_key]

result = np.zeros(original_shape, dtype=np.float32)

box_start = meta_dict[self.start_coord_key]

box_end = meta_dict[self.end_coord_key]

spatial_dims = min(len(box_start), len(image.shape[1:]))

slices = tuple(

[slice(None)] + [slice(s, e) for s, e in zip(box_start[:spatial_dims], box_end[:spatial_dims])]

)

result[slices] = image

else:

result = image

# Undo Spacing

if self.restore_spacing:

current_size = result.shape[1:]

# change spatial_shape from HWD to DHW

spatial_shape = list(np.roll(meta_dict["spatial_shape"], 1))

spatial_size = spatial_shape[-len(current_size) :]

if np.any(np.not_equal(current_size, spatial_size)):

resizer = Resize(spatial_size=spatial_size, mode=mode)

result = resizer(result, mode=mode, align_corners=align_corners) # type: ignore

# Undo Slicing

slice_idx = meta_dict.get("slice_idx")

final_result: NdarrayOrTensor

if not self.restore_slicing: # do nothing if restore slicing isn't requested

final_result = result

elif slice_idx is None or self.slice_only:

final_result = result if len(result.shape) <= 3 else result[0]

else:

slice_idx = meta_dict["slice_idx"][0]

final_result = np.zeros(tuple(spatial_shape))

final_result[slice_idx] = result

d[key] = final_result

meta_key = meta_key or f"{key}_{self.meta_key_postfix}"

meta = d.get(meta_key)

if meta is None:

meta = dict()

d[meta_key] = meta

meta["slice_idx"] = slice_idx

meta["affine"] = meta_dict["original_affine"]

return d

class Fetch2DSliced(MapTransform):

"""

Fetch one slice in case of a 3D volume.

The volume only contains spatial coordinates.

Args:

keys: keys of the corresponding items to be transformed.

guidance: key that represents guidance.

axis: axis that represents slice in 3D volume.

meta_keys: explicitly indicate the key of the corresponding metadata dictionary.

for example, for data with key `image`, the metadata by default is in `image_meta_dict`.

the metadata is a dictionary object which contains: filename, original_shape, etc.

it can be a sequence of string, map to the `keys`.

if None, will try to construct meta_keys by `key_{meta_key_postfix}`.

meta_key_postfix: use `key_{meta_key_postfix}` to fetch the metadata according to the key data,

default is `meta_dict`, the metadata is a dictionary object.

For example, to handle key `image`, read/write affine matrices from the

metadata `image_meta_dict` dictionary's `affine` field.

allow_missing_keys: don't raise exception if key is missing.

"""

def __init__(

self,

keys: KeysCollection,

guidance: str = "guidance",

axis: int = 0,

meta_keys: KeysCollection | None = None,

meta_key_postfix: str = DEFAULT_POST_FIX,

allow_missing_keys: bool = False,

):

super().__init__(keys, allow_missing_keys)

self.guidance = guidance

self.axis = axis

self.meta_keys = ensure_tuple_rep(None, len(self.keys)) if meta_keys is None else ensure_tuple(meta_keys)

if len(self.keys) != len(self.meta_keys):

raise ValueError("meta_keys should have the same length as keys.")

self.meta_key_postfix = ensure_tuple_rep(meta_key_postfix, len(self.keys))

def _apply(self, image, guidance):

slice_idx = guidance[2] # (pos, neg, slice_idx)

idx = []

for i, size_i in enumerate(image.shape):

idx.append(slice_idx) if i == self.axis else idx.append(slice(0, size_i))

return image[tuple(idx)], tuple(idx)

def __call__(self, data):

d = dict(data)

guidance = d[self.guidance]

if len(guidance) < 3:

raise RuntimeError("Guidance does not container slice_idx!")

for key, meta_key, meta_key_postfix in self.key_iterator(d, self.meta_keys, self.meta_key_postfix):

img_slice, idx = self._apply(d[key], guidance)

d[key] = img_slice

d[meta_key or f"{key}_{meta_key_postfix}"]["slice_idx"] = idx

return d