from typing import (TYPE_CHECKING, Callable, Dict, List, Literal, Optional,
Sequence, Tuple, Union)
from pydantic import PositiveInt as PosInt, conint
import math
import random
import numpy as np
from shapely.geometry import Polygon
from shapely.ops import unary_union
from rasterio.windows import Window as RioWindow
from rastervision.pipeline.utils import repr_with_args
NonNegInt = conint(ge=0)
if TYPE_CHECKING:
from shapely.geometry import MultiPolygon
[docs]class BoxSizeError(ValueError):
pass
[docs]class Box():
"""A multi-purpose box (ie. rectangle) representation."""
[docs] def __init__(self, ymin, xmin, ymax, xmax):
"""Construct a bounding box.
Unless otherwise stated, the convention is that these coordinates are
in pixel coordinates and represent boxes that lie within a
RasterSource.
Args:
ymin: minimum y value (y is row)
xmin: minimum x value (x is column)
ymax: maximum y value
xmax: maximum x value
"""
self.ymin = ymin
self.xmin = xmin
self.ymax = ymax
self.xmax = xmax
def __eq__(self, other: 'Box') -> bool:
"""Return true if other has same coordinates."""
return self.tuple_format() == other.tuple_format()
def __ne__(self, other: 'Box'):
"""Return true if other has different coordinates."""
return self.tuple_format() != other.tuple_format()
@property
def height(self) -> int:
"""Return height of Box."""
return self.ymax - self.ymin
@property
def width(self) -> int:
"""Return width of Box."""
return self.xmax - self.xmin
@property
def extent(self) -> 'Box':
"""Return a (0, 0, h, w) Box representing the size of this Box."""
return Box(0, 0, self.height, self.width)
@property
def size(self) -> Tuple[int, int]:
return self.height, self.width
@property
def area(self) -> int:
"""Return area of Box."""
return self.height * self.width
[docs] def normalize(self) -> 'Box':
"""Ensure ymin <= ymax and xmin <= xmax."""
ymin, ymax = sorted((self.ymin, self.ymax))
xmin, xmax = sorted((self.xmin, self.xmax))
return Box(ymin, xmin, ymax, xmax)
[docs] def to_int(self):
return Box(
int(self.ymin), int(self.xmin), int(self.ymax), int(self.xmax))
[docs] @staticmethod
def to_npboxes(boxes):
"""Return nx4 numpy array from list of Box."""
nb_boxes = len(boxes)
npboxes = np.empty((nb_boxes, 4))
for boxind, box in enumerate(boxes):
npboxes[boxind, :] = box.npbox_format()
return npboxes
def __iter__(self):
return iter(self.tuple_format())
[docs] def __getitem__(self, i):
return self.tuple_format()[i]
def __repr__(self) -> str:
return repr_with_args(self, **self.to_dict())
def __hash__(self) -> int:
return hash(self.tuple_format())
[docs] def geojson_coordinates(self) -> List[Tuple[int, int]]:
"""Return Box as GeoJSON coordinates."""
# Compass directions:
nw = [self.xmin, self.ymin]
ne = [self.xmin, self.ymax]
se = [self.xmax, self.ymax]
sw = [self.xmax, self.ymin]
return [nw, ne, se, sw, nw]
[docs] def make_random_square_container(self, size: int) -> 'Box':
"""Return a new square Box that contains this Box.
Args:
size: the width and height of the new Box
"""
return self.make_random_box_container(size, size)
[docs] def make_random_box_container(self, out_h: int, out_w: int) -> 'Box':
"""Return a new rectangular Box that contains this Box.
Args:
out_h (int): the height of the new Box
out_w (int): the width of the new Box
"""
self_h, self_w = self.size
if out_h < self_h:
raise BoxSizeError('size of random container cannot be < height')
if out_w < self_w:
raise BoxSizeError('size of random container cannot be < width')
ymin, xmin, _, _ = self.normalize()
lb = ymin - (out_h - self_h)
ub = ymin
out_ymin = random.randint(int(lb), int(ub))
lb = xmin - (out_w - self_w)
ub = xmin
out_xmin = random.randint(int(lb), int(ub))
return Box(out_ymin, out_xmin, out_ymin + out_h, out_xmin + out_w)
[docs] def make_random_square(self, size: int) -> 'Box':
"""Return new randomly positioned square Box that lies inside this Box.
Args:
size: the height and width of the new Box
"""
if size >= self.width:
raise BoxSizeError('size of random square cannot be >= width')
if size >= self.height:
raise BoxSizeError('size of random square cannot be >= height')
ymin, xmin, ymax, xmax = self.normalize()
lb = ymin
ub = ymax - size
rand_y = random.randint(int(lb), int(ub))
lb = xmin
ub = xmax - size
rand_x = random.randint(int(lb), int(ub))
return Box.make_square(rand_y, rand_x, size)
[docs] def intersection(self, other: 'Box') -> 'Box':
"""Return the intersection of this Box and the other.
Args:
other: The box to intersect with this one.
Returns:
The intersection of this box and the other one.
"""
if not self.intersects(other):
return Box(0, 0, 0, 0)
box1 = self.normalize()
box2 = other.normalize()
xmin = max(box1.xmin, box2.xmin)
ymin = max(box1.ymin, box2.ymin)
xmax = min(box1.xmax, box2.xmax)
ymax = min(box1.ymax, box2.ymax)
return Box(xmin=xmin, ymin=ymin, xmax=xmax, ymax=ymax)
[docs] def intersects(self, other: 'Box') -> bool:
box1 = self.normalize()
box2 = other.normalize()
if box1.ymax <= box2.ymin or box1.ymin >= box2.ymax:
return False
if box1.xmax <= box2.xmin or box1.xmin >= box2.xmax:
return False
return True
[docs] @staticmethod
def from_npbox(npbox):
"""Return new Box based on npbox format.
Args:
npbox: Numpy array of form [ymin, xmin, ymax, xmax] with float type
"""
return Box(*npbox)
[docs] @staticmethod
def from_shapely(shape):
"""Instantiate from the bounds of a shapely geometry."""
xmin, ymin, xmax, ymax = shape.bounds
return Box(ymin, xmin, ymax, xmax)
[docs] @classmethod
def from_rasterio(self, rio_window: RioWindow) -> 'Box':
"""Instantiate from a rasterio window."""
yslice, xslice = rio_window.toslices()
return Box(yslice.start, xslice.start, yslice.stop, xslice.stop)
[docs] def to_xywh(self) -> Tuple[int, int, int, int]:
"""Convert to (xmin, ymin, width, height) tuple"""
return (self.xmin, self.ymin, self.width, self.height)
[docs] def to_xyxy(self) -> Tuple[int, int, int, int]:
"""Convert to (xmin, ymin, xmax, ymax) tuple"""
return (self.xmin, self.ymin, self.xmax, self.ymax)
[docs] def to_points(self) -> np.ndarray:
"""Get (x, y) coords of each vertex as a 4x2 numpy array."""
return np.array(self.geojson_coordinates()[:4])
[docs] def to_shapely(self) -> Polygon:
"""Convert to shapely Polygon."""
return Polygon.from_bounds(*self.shapely_format())
[docs] def to_rasterio(self) -> RioWindow:
"""Convert to a Rasterio Window."""
return RioWindow.from_slices(*self.normalize().to_slices())
[docs] def to_slices(self,
h_step: Optional[int] = None,
w_step: Optional[int] = None) -> Tuple[slice, slice]:
"""Convert to slices: ymin:ymax[:h_step], xmin:xmax[:w_step]"""
return slice(self.ymin, self.ymax, h_step), slice(
self.xmin, self.xmax, w_step)
[docs] def translate(self, dy: int, dx: int) -> 'Box':
"""Translate window along y and x axes by the given distances."""
ymin, xmin, ymax, xmax = self
return Box(ymin + dy, xmin + dx, ymax + dy, xmax + dx)
[docs] def to_global_coords(self, bbox: 'Box') -> 'Box':
"""Go from bbox coords to global coords.
E.g., Given a box Box(20, 20, 40, 40) and bbox Box(20, 20, 100, 100),
the box becomes Box(40, 40, 60, 60).
Inverse of Box.to_local_coords().
"""
return self.translate(dy=bbox.ymin, dx=bbox.xmin)
[docs] def to_local_coords(self, bbox: 'Box') -> 'Box':
"""Go from to global coords bbox coords.
E.g., Given a box Box(40, 40, 60, 60) and bbox Box(20, 20, 100, 100),
the box becomes Box(20, 20, 40, 40).
Inverse of Box.to_global_coords().
"""
return self.translate(dy=-bbox.ymin, dx=-bbox.xmin)
[docs] def reproject(self, transform_fn: Callable) -> 'Box':
"""Reprojects this box based on a transform function.
Args:
transform_fn: A function that takes in a tuple (x, y) and
reprojects that point to the target coordinate reference
system.
"""
(xmin, ymin) = transform_fn((self.xmin, self.ymin))
(xmax, ymax) = transform_fn((self.xmax, self.ymax))
return Box(ymin, xmin, ymax, xmax)
[docs] @staticmethod
def make_square(ymin, xmin, size) -> 'Box':
"""Return new square Box."""
return Box(ymin, xmin, ymin + size, xmin + size)
[docs] def center_crop(self, edge_offset_y: int, edge_offset_x: int) -> 'Box':
"""Return Box whose sides are eroded by the given offsets.
Box(0, 0, 10, 10).center_crop(2, 4) == Box(2, 4, 8, 6)
"""
return Box(self.ymin + edge_offset_y, self.xmin + edge_offset_x,
self.ymax - edge_offset_y, self.xmax - edge_offset_x)
[docs] def erode(self, erosion_sz) -> 'Box':
"""Return new Box whose sides are eroded by erosion_sz."""
return self.center_crop(erosion_sz, erosion_sz)
[docs] def buffer(self, buffer_sz: float, max_extent: 'Box') -> 'Box':
"""Return new Box whose sides are buffered by buffer_sz.
The resulting box is clipped so that the values of the corners are
always greater than zero and less than the height and width of
max_extent.
"""
buffer_sz = max(0., buffer_sz)
if buffer_sz < 1.:
delta_width = int(round(buffer_sz * self.width))
delta_height = int(round(buffer_sz * self.height))
else:
delta_height = delta_width = int(round(buffer_sz))
return Box(
max(0, math.floor(self.ymin - delta_height)),
max(0, math.floor(self.xmin - delta_width)),
min(max_extent.height,
int(self.ymax) + delta_height),
min(max_extent.width,
int(self.xmax) + delta_width))
[docs] def pad(self, ymin: int, xmin: int, ymax: int, xmax: int) -> 'Box':
"""Pad sides by the given amount."""
return Box(
ymin=self.ymin - ymin,
xmin=self.xmin - xmin,
ymax=self.ymax + ymax,
xmax=self.xmax + xmax)
[docs] def copy(self) -> 'Box':
return Box(*self)
[docs] def get_windows(self,
size: Union[PosInt, Tuple[PosInt, PosInt]],
stride: Union[PosInt, Tuple[PosInt, PosInt]],
padding: Optional[Union[NonNegInt, Tuple[
NonNegInt, NonNegInt]]] = None,
pad_direction: Literal['both', 'start', 'end'] = 'end'
) -> List['Box']:
"""Returns a list of boxes representing windows generated using a
sliding window traversal with the specified size, stride, and
padding.
Each of size, stride, and padding can be either a positive int or
a tuple `(vertical-component, horizontal-component)` of positive ints.
Padding currently only applies to the right and bottom edges.
Args:
size (Union[PosInt, Tuple[PosInt, PosInt]]): Size (h, w) of the
windows.
stride (Union[PosInt, Tuple[PosInt, PosInt]]): Step size between
windows. Can be 2-tuple (h_step, w_step) or positive int.
padding (Optional[Union[PosInt, Tuple[PosInt, PosInt]]], optional):
Optional padding to accommodate windows that overflow the
extent. Can be 2-tuple (h_pad, w_pad) or non-negative int.
If None, will be set to (size[0]//2, size[1]//2).
Defaults to None.
pad_direction (Literal['both', 'start', 'end']): If 'end', only pad
ymax and xmax (bottom and right). If 'start', only pad ymin and
xmin (top and left). If 'both', pad all sides. Has no effect if
padding is zero. Defaults to 'end'.
Returns:
List[Box]: List of Box objects.
"""
if not isinstance(size, tuple):
size = (size, size)
if not isinstance(stride, tuple):
stride = (stride, stride)
if size[0] <= 0 or size[1] <= 0 or stride[0] <= 0 or stride[1] <= 0:
raise ValueError('size and stride must be positive.')
if padding is None:
padding = (size[0] // 2, size[1] // 2)
if not isinstance(padding, tuple):
padding = (padding, padding)
if padding[0] < 0 or padding[1] < 0:
raise ValueError('padding must be non-negative.')
if padding != (0, 0):
h_pad, w_pad = padding
if pad_direction == 'both':
padded_box = self.pad(
ymin=h_pad, xmin=w_pad, ymax=h_pad, xmax=w_pad)
elif pad_direction == 'end':
padded_box = self.pad(ymin=0, xmin=0, ymax=h_pad, xmax=w_pad)
elif pad_direction == 'start':
padded_box = self.pad(ymin=h_pad, xmin=w_pad, ymax=0, xmax=0)
else:
raise ValueError('pad_directions must be one of: '
'"both", "start", "end".')
return padded_box.get_windows(
size=size, stride=stride, padding=(0, 0))
# padding is necessarily (0, 0) at this point, so we ignore it
h, w = size
h_step, w_step = stride
# lb = lower bound, ub = upper bound
ymin_lb = self.ymin
xmin_lb = self.xmin
ymin_ub = self.ymax - h
xmin_ub = self.xmax - w
windows = []
for ymin in range(ymin_lb, ymin_ub + 1, h_step):
for xmin in range(xmin_lb, xmin_ub + 1, w_step):
windows.append(Box(ymin, xmin, ymin + h, xmin + w))
return windows
[docs] def to_dict(self) -> Dict[str, int]:
"""Convert to a dict with keys: ymin, xmin, ymax, xmax."""
return {
'ymin': self.ymin,
'xmin': self.xmin,
'ymax': self.ymax,
'xmax': self.xmax,
}
[docs] @classmethod
def from_dict(cls, d: dict) -> 'Box':
return cls(d['ymin'], d['xmin'], d['ymax'], d['xmax'])
[docs] @staticmethod
def filter_by_aoi(windows: List['Box'],
aoi_polygons: List[Polygon],
within: bool = True) -> List['Box']:
"""Filters windows by a list of AOI polygons
Args:
within: if True, windows are only kept if they lie fully within an
AOI polygon. Otherwise, windows are kept if they intersect an
AOI polygon.
"""
# merge overlapping polygons, if any
aoi_polygons: Polygon | MultiPolygon = unary_union(aoi_polygons)
if within:
keep_window = aoi_polygons.contains
else:
keep_window = aoi_polygons.intersects
out = [w for w in windows if keep_window(w.to_shapely())]
return out
[docs] @staticmethod
def within_aoi(window: 'Box',
aoi_polygons: Polygon | List[Polygon]) -> bool:
"""Check if window is within the union of given AOI polygons."""
aoi_polygons: Polygon | MultiPolygon = unary_union(aoi_polygons)
w = window.to_shapely()
out = aoi_polygons.contains(w)
return out
[docs] @staticmethod
def intersects_aoi(window: 'Box',
aoi_polygons: Polygon | List[Polygon]) -> bool:
"""Check if window intersects with the union of given AOI polygons."""
aoi_polygons: Polygon | MultiPolygon = unary_union(aoi_polygons)
w = window.to_shapely()
out = aoi_polygons.intersects(w)
return out
[docs] def __contains__(self, query: Union['Box', Sequence]) -> bool:
"""Check if box or point is contained within this box.
Args:
query: Box or single point (x, y).
Raises:
NotImplementedError: if query is not a Box or tuple/list.
"""
if isinstance(query, Box):
ymin, xmin, ymax, xmax = query
return (ymin >= self.ymin and xmin >= self.xmin
and ymax <= self.ymax and xmax <= self.xmax)
elif isinstance(query, (tuple, list)):
x, y = query
return self.xmin <= x <= self.xmax and self.ymin <= y <= self.ymax
else:
raise NotImplementedError()