249 lines
8.6 KiB
Python
249 lines
8.6 KiB
Python
from typing import Any, cast, Self
|
|
from collections.abc import Iterator
|
|
import copy
|
|
import functools
|
|
|
|
import numpy
|
|
from numpy import pi
|
|
from numpy.typing import NDArray, ArrayLike
|
|
|
|
from . import Shape, normalized_shape_tuple
|
|
from .polygon import Polygon
|
|
from ..error import PatternError
|
|
from ..repetition import Repetition
|
|
from ..utils import annotations_lt, annotations_eq, rep2key, annotations_t
|
|
|
|
|
|
def _normalize_rects(rects: ArrayLike) -> NDArray[numpy.float64]:
|
|
arr = numpy.asarray(rects, dtype=float)
|
|
if arr.ndim != 2 or arr.shape[1] != 4:
|
|
raise PatternError('Rectangles must be an Nx4 array of [xmin, ymin, xmax, ymax]')
|
|
if numpy.any(arr[:, 0] > arr[:, 2]) or numpy.any(arr[:, 1] > arr[:, 3]):
|
|
raise PatternError('Rectangles must satisfy xmin <= xmax and ymin <= ymax')
|
|
if arr.shape[0] <= 1:
|
|
return arr
|
|
order = numpy.lexsort((arr[:, 3], arr[:, 2], arr[:, 1], arr[:, 0]))
|
|
return arr[order]
|
|
|
|
|
|
def _renormalize_rects_in_place(rects: NDArray[numpy.float64]) -> None:
|
|
x0 = numpy.minimum(rects[:, 0], rects[:, 2])
|
|
x1 = numpy.maximum(rects[:, 0], rects[:, 2])
|
|
y0 = numpy.minimum(rects[:, 1], rects[:, 3])
|
|
y1 = numpy.maximum(rects[:, 1], rects[:, 3])
|
|
rects[:, 0] = x0
|
|
rects[:, 1] = y0
|
|
rects[:, 2] = x1
|
|
rects[:, 3] = y1
|
|
|
|
|
|
@functools.total_ordering
|
|
class RectCollection(Shape):
|
|
"""
|
|
A collection of axis-aligned rectangles, stored as an Nx4 array of
|
|
`[xmin, ymin, xmax, ymax]` rows.
|
|
"""
|
|
__slots__ = (
|
|
'_rects',
|
|
'_repetition', '_annotations',
|
|
)
|
|
|
|
_rects: NDArray[numpy.float64]
|
|
|
|
@property
|
|
def rects(self) -> NDArray[numpy.float64]:
|
|
return self._rects
|
|
|
|
@rects.setter
|
|
def rects(self, val: ArrayLike) -> None:
|
|
self._rects = _normalize_rects(val)
|
|
|
|
@property
|
|
def offset(self) -> NDArray[numpy.float64]:
|
|
return numpy.zeros(2)
|
|
|
|
@offset.setter
|
|
def offset(self, val: ArrayLike) -> None:
|
|
if numpy.any(val):
|
|
raise PatternError('RectCollection offset is forced to (0, 0)')
|
|
|
|
def set_offset(self, val: ArrayLike) -> Self:
|
|
if numpy.any(val):
|
|
raise PatternError('RectCollection offset is forced to (0, 0)')
|
|
return self
|
|
|
|
def translate(self, offset: ArrayLike) -> Self:
|
|
delta = numpy.asarray(offset, dtype=float).reshape(2)
|
|
self._rects[:, [0, 2]] += delta[0]
|
|
self._rects[:, [1, 3]] += delta[1]
|
|
return self
|
|
|
|
def __init__(
|
|
self,
|
|
rects: ArrayLike,
|
|
*,
|
|
offset: ArrayLike = (0.0, 0.0),
|
|
rotation: float = 0.0,
|
|
repetition: Repetition | None = None,
|
|
annotations: annotations_t = None,
|
|
) -> None:
|
|
self.rects = rects
|
|
self.repetition = repetition
|
|
self.annotations = annotations
|
|
if rotation:
|
|
self.rotate(rotation)
|
|
if numpy.any(offset):
|
|
self.translate(offset)
|
|
|
|
@classmethod
|
|
def _from_raw(
|
|
cls,
|
|
*,
|
|
rects: NDArray[numpy.float64],
|
|
annotations: annotations_t = None,
|
|
repetition: Repetition | None = None,
|
|
) -> Self:
|
|
new = cls.__new__(cls)
|
|
new._rects = rects
|
|
new._repetition = repetition
|
|
new._annotations = annotations
|
|
return new
|
|
|
|
@property
|
|
def polygon_vertices(self) -> Iterator[NDArray[numpy.float64]]:
|
|
for rect in self._rects:
|
|
xmin, ymin, xmax, ymax = rect
|
|
yield numpy.array([
|
|
[xmin, ymin],
|
|
[xmin, ymax],
|
|
[xmax, ymax],
|
|
[xmax, ymin],
|
|
], dtype=float)
|
|
|
|
def __deepcopy__(self, memo: dict | None = None) -> Self:
|
|
memo = {} if memo is None else memo
|
|
new = copy.copy(self)
|
|
new._rects = self._rects.copy()
|
|
new._repetition = copy.deepcopy(self._repetition, memo)
|
|
new._annotations = copy.deepcopy(self._annotations)
|
|
return new
|
|
|
|
def _sorted_rects(self) -> NDArray[numpy.float64]:
|
|
if self._rects.shape[0] <= 1:
|
|
return self._rects
|
|
order = numpy.lexsort((self._rects[:, 3], self._rects[:, 2], self._rects[:, 1], self._rects[:, 0]))
|
|
return self._rects[order]
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
return (
|
|
type(self) is type(other)
|
|
and numpy.array_equal(self._sorted_rects(), other._sorted_rects())
|
|
and self.repetition == other.repetition
|
|
and annotations_eq(self.annotations, other.annotations)
|
|
)
|
|
|
|
def __lt__(self, other: Shape) -> bool:
|
|
if type(self) is not type(other):
|
|
if repr(type(self)) != repr(type(other)):
|
|
return repr(type(self)) < repr(type(other))
|
|
return id(type(self)) < id(type(other))
|
|
|
|
other = cast('RectCollection', other)
|
|
self_rects = self._sorted_rects()
|
|
other_rects = other._sorted_rects()
|
|
if not numpy.array_equal(self_rects, other_rects):
|
|
min_len = min(self_rects.shape[0], other_rects.shape[0])
|
|
eq_mask = self_rects[:min_len] != other_rects[:min_len]
|
|
eq_lt = self_rects[:min_len] < other_rects[:min_len]
|
|
eq_lt_masked = eq_lt[eq_mask]
|
|
if eq_lt_masked.size > 0:
|
|
return bool(eq_lt_masked.flat[0])
|
|
return self_rects.shape[0] < other_rects.shape[0]
|
|
if self.repetition != other.repetition:
|
|
return rep2key(self.repetition) < rep2key(other.repetition)
|
|
return annotations_lt(self.annotations, other.annotations)
|
|
|
|
def to_polygons(
|
|
self,
|
|
num_vertices: int | None = None, # unused # noqa: ARG002
|
|
max_arclen: float | None = None, # unused # noqa: ARG002
|
|
) -> list[Polygon]:
|
|
return [
|
|
Polygon(
|
|
vertices=vertices,
|
|
repetition=copy.deepcopy(self.repetition),
|
|
annotations=copy.deepcopy(self.annotations),
|
|
)
|
|
for vertices in self.polygon_vertices
|
|
]
|
|
|
|
def get_bounds_single(self) -> NDArray[numpy.float64] | None:
|
|
if self._rects.size == 0:
|
|
return None
|
|
mins = self._rects[:, :2].min(axis=0)
|
|
maxs = self._rects[:, 2:].max(axis=0)
|
|
return numpy.vstack((mins, maxs))
|
|
|
|
def rotate(self, theta: float) -> Self:
|
|
quarter_turns = int(numpy.rint(theta / (pi / 2)))
|
|
if not numpy.isclose(theta, quarter_turns * (pi / 2)):
|
|
raise PatternError('RectCollection only supports Manhattan rotations')
|
|
turns = quarter_turns % 4
|
|
if turns == 0 or self._rects.size == 0:
|
|
return self
|
|
|
|
corners = numpy.stack((
|
|
self._rects[:, [0, 1]],
|
|
self._rects[:, [0, 3]],
|
|
self._rects[:, [2, 3]],
|
|
self._rects[:, [2, 1]],
|
|
), axis=1)
|
|
flat = corners.reshape(-1, 2)
|
|
if turns == 1:
|
|
rotated = numpy.column_stack((-flat[:, 1], flat[:, 0]))
|
|
elif turns == 2:
|
|
rotated = -flat
|
|
else:
|
|
rotated = numpy.column_stack((flat[:, 1], -flat[:, 0]))
|
|
corners = rotated.reshape(corners.shape)
|
|
self._rects[:, 0] = corners[:, :, 0].min(axis=1)
|
|
self._rects[:, 1] = corners[:, :, 1].min(axis=1)
|
|
self._rects[:, 2] = corners[:, :, 0].max(axis=1)
|
|
self._rects[:, 3] = corners[:, :, 1].max(axis=1)
|
|
return self
|
|
|
|
def mirror(self, axis: int = 0) -> Self:
|
|
if axis not in (0, 1):
|
|
raise PatternError('Axis must be 0 or 1')
|
|
if axis == 0:
|
|
self._rects[:, [1, 3]] *= -1
|
|
else:
|
|
self._rects[:, [0, 2]] *= -1
|
|
_renormalize_rects_in_place(self._rects)
|
|
return self
|
|
|
|
def scale_by(self, c: float) -> Self:
|
|
self._rects *= c
|
|
_renormalize_rects_in_place(self._rects)
|
|
return self
|
|
|
|
def normalized_form(self, norm_value: float) -> normalized_shape_tuple:
|
|
rects = self._sorted_rects()
|
|
centers = 0.5 * (rects[:, :2] + rects[:, 2:])
|
|
offset = centers.mean(axis=0)
|
|
zeroed = rects.copy()
|
|
zeroed[:, [0, 2]] -= offset[0]
|
|
zeroed[:, [1, 3]] -= offset[1]
|
|
normed = zeroed / norm_value
|
|
return (
|
|
(type(self), normed.data.tobytes()),
|
|
(offset, 1.0, 0.0, False),
|
|
lambda: RectCollection(rects=normed * norm_value),
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
if self._rects.size == 0:
|
|
return '<RectCollection r0>'
|
|
centers = 0.5 * (self._rects[:, :2] + self._rects[:, 2:])
|
|
centroid = centers.mean(axis=0)
|
|
return f'<RectCollection centroid {centroid} r{self._rects.shape[0]}>'
|