# -*- encoding: utf-8 -*-
# @File : shapes.py
# @Time : 2019-11-20
# @Author : zjh
r"""
Common shapes for object detection
"""
__all__ = ["Shape", "Point", "Points", "Line", "Polygon", "Box", "Boxes", "Mask"]
import numpy as np
import cv2
import re
class _ORDER:
width_height = 0
height_width = 1
def _check_shape_compatible(template, shape):
exp_shape = [x.strip() for x in template.strip("()").split(",")]
if len(exp_shape) != len(shape):
return False
for e, s in zip(exp_shape, shape):
if e == "?":
continue
elif e.startswith(">") or e.startswith("<"):
if not eval(str(s) + e):
return False
else:
if not (int(e) == s):
return False
return True
def _make_dim_compatible(template, obj, force=False):
exp_shape = [x.strip() for x in template.strip("()").split(",")]
shape = [int(x) if re.match(r"^[1-9]\d*$", x) else 0 for x in exp_shape]
if np.product(shape) == 0 or force:
obj = np.reshape(obj, shape)
return obj
[docs]class Shape:
"""
Base class for object detection shapes. Wrap a numpy.array object and
add specified operations for each type shape.
"""
order = _ORDER.width_height
_shape = ""
def __init__(self, obj, dtype=None):
obj = np.asarray(obj, dtype=dtype)
# set expect dim for empty object
if len(obj) == 0:
obj = _make_dim_compatible(self._shape, obj)
if not _check_shape_compatible(self._shape, obj.shape):
raise ValueError("Expect array shape %s, but get %s" % (self._shape, obj.shape))
self._obj = obj
def __getattr__(self, item):
return getattr(self._obj, item)
def __repr__(self):
return "{}({!r})".format(self.__class__.__name__, self._obj)
def __len__(self):
return len(self._obj)
def __getitem__(self, *args, **kwargs):
return self._obj.__getitem__(*args, **kwargs)
def __setitem__(self, *args, **kwargs):
self._obj.__setitem__(*args, **kwargs)
def __delitem__(self, *args, **kwargs):
self._obj.__delitem__(*args, **kwargs)
def __add__(self, other):
return self._obj.__add__(other)
def __sub__(self, other):
return self._obj.__sub__(other)
def __mul__(self, other):
return self._obj.__mul__(other)
def __truediv__(self, other):
return self._obj.__truediv__(other)
def __mod__(self, other):
return self._obj.__mod__(other)
[docs] def numpy(self):
"""
Get numpy data
:return: shape data
:rtype: np.ndarray
"""
return self._obj
[docs] def swap(self):
"""
Swap this shape coordinate order from x-y to y-x
:return: swapped this shape
:rtype: Shape
"""
shape = self._obj.shape
obj = np.reshape(self._obj, [-1, 2])
obj = obj[..., ::-1]
self._obj = np.reshape(obj, shape)
return self
[docs]class Point(Shape):
"""
Point of 2d
"""
_shape = "(2)"
[docs]class Points(Point):
"""
Collection of Points
"""
_shape = "(?, 2)"
[docs] def bounding_box(self):
"""
Get the minimum bounding box a collection of points
:return: bounding box
:rtype: Box
"""
if len(self._obj) == 0:
raise ValueError("Empty points have no bounding box.")
top_left = np.min(self._obj, axis=0)
bot_right = np.max(self._obj, axis=0)
return Box(np.concatenate([top_left, bot_right]))
[docs]class Line(Points):
"""
Line of 2d
"""
[docs] def length(self):
"""
Calculate the length of a continuous line
:return: the line length
:rtype: np.float
"""
if len(self._obj) < 2:
return 0.
dis = self._obj[:-1] - self._obj[1:]
lth = np.sqrt(np.sum(dis ** 2, axis=-1)).sum()
return lth
[docs]class Box(Shape):
"""
Box of 2d. Record top_left and bottom_right corner position of box.
"""
_shape = "(4)"
[docs] def area(self):
"""
Calculate box area
:return: box area
:rtype: np.float
"""
return (self._obj[0] - self._obj[2]) * (self._obj[1] - self._obj[3])
[docs] def to_polygon(self):
"""
Convert box to polygon from top_left and across top_right, bot_right
and end to bot_left
:return: the converted polygon
:rtype: Polygon
"""
top_left = self._obj[:2]
top_right = [self._obj[0], self._obj[3]]
bot_left = [self._obj[2], self._obj[1]]
bot_right = self._obj[2:]
return Polygon([top_left, top_right, bot_right, bot_left])
[docs] def to_mask(self, size=None):
"""
Convert box to mask
:param size:
:return:
"""
return self.to_polygon().to_mask(size)
[docs] def bsize(self):
"""
Size of box
:return: box size in format np.array([width, height])
:rtype: np.ndarray
"""
return self._obj[2:] - self._obj[:2]
[docs] def center(self):
"""
Center point of a box
:return: center point
:rtype: Point
"""
return Point([(self._obj[2] + self._obj[0]) / 2,
(self._obj[3] + self._obj[1]) / 2])
[docs] def to_cxywh(self):
"""
Convert box format to [center-x, center-y, width, height]
:return: converted box in format [center-x, center-y, width, height]
:rtype: np.ndarray
"""
return np.asarray(tuple(self.center()) + tuple(self.bsize()))
[docs] @classmethod
def from_cxywh(cls, cxywh):
"""
Create box from format [min-x, min-y, max-x, max-y]
:return: converted box in format [min-x, min-y, max-x, max-y]
:rtype: Box
"""
cx, cy, w, h = cxywh
tx, ty = cx - w / 2, cy - h / 2
bx, by = cx + w / 2, cy + h / 2
return cls([tx, ty, bx, by])
[docs] def scale(self, scale):
"""
Scale box
:param scale: scale factor
:return: a scaled box
:rtype: Box
"""
cxywh = self.to_cxywh()
cxywh[2:] *= scale
return self.from_cxywh(cxywh)
[docs]class Boxes(Shape):
"""Collection of Box"""
_shape = "(?, 4)"
[docs] def areas(self):
"""
Calculate boxes areas
:return: boxes areas
:rtype: np.ndarray
"""
return (self._obj[:, 0] - self._obj[:, 2]) * (self._obj[:, 1] - self._obj[:, 3])
[docs] def bsize(self):
"""
Sizes of boxes
:return: boxes sizes in format np.array([[width, height], ...])
:rtype: np.ndarray
"""
return self._obj[:, 2:] - self._obj[:, :2]
[docs] def center(self):
"""
Center points of a boxes
:return: center points
:rtype: Points
"""
return Points((self._obj[:, :2] + self._obj[:, 2:]) / 2)
[docs] def to_cxywh(self):
"""
Convert box format to [center-x, center-y, width, height]
:return: converted boxes in format [center-x, center-y, width, height]
:rtype: np.ndarray
"""
return np.concatenate([self.center(), self.bsize()], axis=1)
[docs] @classmethod
def from_cxywh(cls, cxywh):
"""
Create box from format [min-x, min-y, max-x, max-y]
:return: converted box in format [min-x, min-y, max-x, max-y]
:rtype: Boxes
"""
tx, ty = cxywh[:, 0] - cxywh[:, 2] / 2, cxywh[:, 1] - cxywh[:, 3] / 2
bx, by = cxywh[:, 0] + cxywh[:, 2] / 2, cxywh[:, 1] + cxywh[:, 3] / 2
return cls(np.stack([tx, ty, bx, by], axis=1))
[docs] def scale(self, scale):
"""
Scale boxes
:param scale: scale factor
:return: a scaled boxes
:rtype: Boxes
"""
cxywh = self.to_cxywh()
cxywh[:, 2:] *= scale
return self.from_cxywh(cxywh)
[docs]class Polygon(Shape):
"""
Polygon of 2d
"""
_shape = "(>=3, 2)"
[docs] def area(self):
"""
Calculate polygon area use mask area calculate
:return: polygon area
:rtype: np.float
"""
top_left = np.min(self._obj, axis=0)
mask = Polygon(self._obj - top_left).to_mask()
return mask.area()
[docs] def bounding_box(self):
"""
Get the minimum bounding box of this polygon
:return: bounding box
:rtype: Box
"""
return Points(self._obj).bounding_box()
[docs] def to_mask(self, size=None):
"""
Convert polygon to mask
:param size: the final mask size. Default is ``None`` means use the minimum
size that can overlap this mask
:return: converted mask from polygon
:rtype: np.ndarray
"""
if size is None:
bot_right = np.max(self._obj, axis=0)
size = [int(np.ceil(x)) for x in bot_right]
if self.order == _ORDER.width_height:
size = list(reversed(size))
else:
if len(size) != 2:
raise ValueError("mask size must be tuple/list like (height, width)")
mask = np.zeros(size, dtype=np.uint8)
if self.order == _ORDER.width_height:
mask = cv2.fillPoly(mask, [self._obj], (255,))
else:
mask = cv2.fillPoly(mask, [self._obj[:, ::-1]], (255,))
return Mask(mask)
[docs]class Mask(Shape):
"""
Mask of 2d
"""
_shape = "(?, ?)"
[docs] def area(self):
"""
Calculate mask area
:return: mask area
:rtype: np.float
"""
return np.not_equal(self._obj, 0).sum()
[docs] def bounding_box(self):
"""
Get the minimum bounding box of this mask
:return: then minimum bounding box of this mask
:rtype: Box
"""
idx0, idx1 = np.where(self._obj)
if len(idx0) == 0:
return Box([0, 0, 0, 0])
if self.order == _ORDER.width_height:
points = np.stack([idx1, idx0], axis=-1)
else:
points = np.stack([idx0, idx1], axis=-1)
return Points(points).bounding_box()
[docs] def swap(self):
"""
Swap mask shape coordinate order from x-y to y-x
:return: swapped this mask
:rtype: Mask
"""
self._obj = self._obj.T
return self