Source code for dpipe.im.box

"""
Functions to work with boxes: immutable numpy arrays of shape (2, n) which represent
the coordinates of the upper left and lower right corners of an n-dimensional rectangle.

In slicing operations, as everywhere in Python, the left corner is inclusive, and the right one is non-inclusive.
"""
import itertools
from functools import wraps
from typing import Callable

import numpy as np

from ..checks import check_len
from .shape_utils import compute_shape_from_spatial
from .utils import build_slices

# box type
Box = np.ndarray


[docs]def make_box_(iterable) -> Box: """ Returns a box, generated inplace from the ``iterable``. If ``iterable`` was a numpy array, will make it immutable and return. """ box = np.asarray(iterable) box.setflags(write=False) assert box.ndim == 2 and len(box) == 2, box.shape assert np.all(box[0] <= box[1]), box return box
def get_volume(box: Box): return np.prod(box[1] - box[0], axis=0)
[docs]def returns_box(func: Callable) -> Callable: """Returns function, decorated so that it returns a box.""" @wraps(func) def func_returning_box(*args, **kwargs): return make_box_(func(*args, **kwargs)) func_returning_box.__annotations__['return'] = Box return func_returning_box
[docs]@returns_box def get_containing_box(shape: tuple): """Returns box that contains complete array of shape ``shape``.""" return [0] * len(shape), shape
[docs]@returns_box def broadcast_box(box: Box, shape: tuple, dims: tuple): """ Returns box, such that it contains ``box`` across ``dims`` and whole array with shape ``shape`` across other dimensions. """ return (compute_shape_from_spatial([0] * len(shape), box[0], dims), compute_shape_from_spatial(shape, box[1], dims))
[docs]@returns_box def limit_box(box, limit): """ Returns a box, maximum subset of the input ``box`` so that start would be non-negative and stop would be limited by the ``limit``. """ check_len(*box, limit) return np.maximum(box[0], 0), np.minimum(box[1], limit)
[docs]def get_box_padding(box: Box, limit): """ Returns padding that is necessary to get ``box`` from array of shape ``limit``. Returns padding in numpy form, so it can be given to `numpy.pad`. """ check_len(*box, limit) return np.maximum([-box[0], box[1] - limit], 0).T
@returns_box def get_union_box(*boxes): start = np.min([box[0] for box in boxes], axis=0) stop = np.max([box[1] for box in boxes], axis=0) return start, stop
[docs]@returns_box def add_margin(box: Box, margin): """ Returns a box with size increased by the ``margin`` (need to be broadcastable to the box) compared to the input ``box``. """ margin = np.broadcast_to(margin, box.shape) return box[0] - margin[0], box[1] + margin[1]
[docs]@returns_box def get_centered_box(center: np.ndarray, box_size: np.ndarray): """ Get box of size ``box_size``, centered in the ``center``. If ``box_size`` is odd, ``center`` will be closer to the right. """ start = center - box_size // 2 stop = center + box_size // 2 + box_size % 2 return start, stop
[docs]@returns_box def mask2bounding_box(mask: np.ndarray): """ Find the smallest box that contains all true values of the ``mask``. """ if not mask.any(): raise ValueError('The mask is empty.') start, stop = [], [] for ax in itertools.combinations(range(mask.ndim), mask.ndim - 1): nonzero = np.any(mask, axis=ax) if np.any(nonzero): left, right = np.where(nonzero)[0][[0, -1]] else: left, right = 0, 0 start.insert(0, left) stop.insert(0, right + 1) return start, stop
def box2slices(box: Box): return build_slices(*box)