def __call__(self, prob_map: NdarrayOrTensor): """ prob_map: the input probabilities map, it must have shape (H[, W, ...]). """ if self.sigma != 0: if not isinstance(prob_map, torch.Tensor): prob_map = torch.as_tensor(prob_map, dtype=torch.float) self.filter.to(prob_map.device) prob_map = self.filter(prob_map) prob_map_shape = prob_map.shape outputs = [] while prob_map.max() > self.prob_threshold: max_idx = unravel_index(prob_map.argmax(), prob_map_shape) prob_max = prob_map[tuple(max_idx)] max_idx = max_idx.cpu().numpy() if isinstance( max_idx, torch.Tensor) else max_idx prob_max = prob_max.item() if isinstance( prob_max, torch.Tensor) else prob_max outputs.append([prob_max] + list(max_idx)) idx_min_range = (max_idx - self.box_lower_bd).clip(0, None) idx_max_range = (max_idx + self.box_upper_bd).clip( None, prob_map_shape) # for each dimension, set values during index ranges to 0 slices = tuple( slice(idx_min_range[i], idx_max_range[i]) for i in range(self.spatial_dims)) prob_map[slices] = 0 return outputs
def __call__(self, img: NdarrayOrTensor, randomize: bool = True) -> NdarrayOrTensor: """ Apply the transform to `img`, if `randomize` randomizing the smooth field otherwise reusing the previous. """ img = convert_to_tensor(img, track_meta=get_track_meta()) if randomize: self.randomize() if not self._do_transform: return img img_min = img.min() img_max = img.max() img_rng = img_max - img_min field = self.sfield() rfield, *_ = convert_to_dst_type(field, img) # everything below here is to be computed using the destination type (numpy, tensor, etc.) img = (img - img_min) / (img_rng + 1e-10) # rescale to unit values img = img**rfield # contrast is changed by raising image data to a power, in this case the field out = (img * img_rng ) + img_min # rescale back to the original image value range return out