refactored dataset's bounding box mixins

cvdatasets/dataset/mixins/bounding_box.py

@@ -1,60 +0,0 @@
-import numpy as np
-
-from cvdatasets.dataset.mixins.base import BaseMixin
-
-class BBoxMixin(BaseMixin):
-
-	def bounding_box(self, i):
-		bbox = self._get("bounding_box", i)
-		return [bbox[attr] for attr in "xywh"]
-
-class MultiBoxMixin(BaseMixin):
-	_all_keys=[
-		"x", "x0", "x1",
-		"y", "y0", "y1",
-		"w", "h",
-	]
-
-	def multi_box(self, i, keys=["x0","x1","y0","y1"]):
-		assert all([key in self._all_keys for key in keys]), \
-			f"unknown keys found: {keys}. Possible are: {self._all_keys}"
-
-		boxes = [
-			dict(
-				x=box["x0"], x0=box["x0"], x1=box["x1"],
-
-				y=box["y0"], y0=box["y0"], y1=box["y1"],
-
-				w=box["x1"] - box["x0"],
-				h=box["y1"] - box["y0"],
-			)
-			for box in self._get("multi_box", i)["objects"]
-		]
-
-		return [[box[key] for key in keys] for box in boxes]
-
-class BBCropMixin(BBoxMixin):
-
-	def __init__(self, *, crop_to_bb=False, crop_uniform=False, **kwargs):
-		super(BBCropMixin, self).__init__(**kwargs)
-		self.crop_to_bb = crop_to_bb
-		self.crop_uniform = crop_uniform
-
-	def bounding_box(self, i):
-		x,y,w,h = super(BBCropMixin, self).bounding_box(i)
-		if self.crop_uniform:
-			x0 = x + w//2
-			y0 = y + h//2
-
-			crop_size = max(w//2, h//2)
-
-			x,y = max(x0 - crop_size, 0), max(y0 - crop_size, 0)
-			w = h = crop_size * 2
-		return x,y,w,h
-
-	def get_example(self, i):
-		im_obj = super(BBCropMixin, self).get_example(i)
-		if self.crop_to_bb:
-			bb = self.bounding_box(i)
-			return im_obj.crop(*bb)
-		return im_obj

cvdatasets/dataset/mixins/bounding_box/__init__.py

@@ -0,0 +1,4 @@
+from cvdatasets.dataset.mixins.bounding_box.base import BBCropMixin
+from cvdatasets.dataset.mixins.bounding_box.base import BBoxMixin
+from cvdatasets.dataset.mixins.bounding_box.bbox import BoundingBox
+from cvdatasets.dataset.mixins.bounding_box.multi_box import MultiBoxMixin

cvdatasets/dataset/mixins/bounding_box/base.py

@@ -0,0 +1,31 @@
+from cvdatasets.dataset.mixins.base import BaseMixin
+from cvdatasets.dataset.mixins.bounding_box.bbox import BoundingBox
+
+class BBoxMixin(BaseMixin):
+
+	def bounding_box(self, i):
+		bbox = self._get("bounding_box", i)
+		return BoundingBox(*bbox)
+
+
+class BBCropMixin(BBoxMixin):
+
+	def __init__(self, *, crop_to_bb=False, crop_uniform=False, **kwargs):
+		super(BBCropMixin, self).__init__(**kwargs)
+		self.crop_to_bb = crop_to_bb
+		self.crop_uniform = crop_uniform
+
+	def bounding_box(self, i):
+		bbox = super(BBCropMixin, self).bounding_box(i)
+
+		if self.crop_uniform:
+			return bbox.squared()
+
+		return bbox
+
+	def get_example(self, i):
+		im_obj = super(BBCropMixin, self).get_example(i)
+		if self.crop_to_bb:
+			bb = self.bounding_box(i)
+			return im_obj.crop(*bb)
+		return im_obj

cvdatasets/dataset/mixins/bounding_box/bbox.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import typing as T
+
+class BoundingBox(T.NamedTuple):
+
+	x: int = 0
+	y: int = 0
+	w: int = 0
+	h: int = 0
+
+	@property
+	def x0(self) -> int:
+		return self.x
+
+	@property
+	def x1(self) -> int:
+		return self.x + self.w
+
+	@property
+	def y0(self) -> int:
+		return self.y
+
+	@property
+	def y1(self) -> int:
+		return self.y + self.h
+
+	@property
+	def center(self) -> T.Tuple[int, int]:
+		return self.x + self.w // 2, self.y + self.h // 2
+
+	def squared(self, *, smaller_size: bool = False) -> BoundingBox:
+		cx, cy = self.center
+		func = min if smaller_size else max
+
+		size = func(self.w, self.h)
+		x0 = max(cx - size // 2, 0)
+		y0 = max(cy - size // 2, 0)
+
+		return BoundingBox(x0, y0, size, size)
+
+	@classmethod
+	def new(cls, **kwargs) -> BoundingBox:
+		# parameters that should be passed once
+		once = [
+			("x", "x0"),
+			("y", "y0"),
+			("w", "x1"),
+			("h", "y1"),
+		]
+
+		for p1, p2 in once:
+			assert (p1 in kwargs and p2 not in kwargs) \
+				or (p2 in kwargs and p1 not in kwargs), \
+				f"please pass only one of these arguments: {p1}, {p2}"
+
+		x = kwargs.get("x", kwargs["x0"])
+		y = kwargs.get("y", kwargs["y0"])
+
+		if "x1" in kwargs:
+			w = kwargs["x1"] - x
+
+		if "y1" in kwargs:
+			h = kwargs["y1"] - y
+
+		return cls(x, y, w, h)
+
+	def get(self, *attrs) -> T.Tuple:
+		return tuple(getattr(self, attr) for attr in attrs)
+
+
+	def __repr__(self) -> str:
+		return f"<BoundingBox [0x{id(self):x}]: (({self.x0}, {self.y0}), ({self.x1}, {self.y1})) [{self.w}x{self.h} px]>"
+
+
+
+if __name__ == '__main__':
+
+	print(BoundingBox.new(x0=1, x1=200, y0=10, y1=90))

cvdatasets/dataset/mixins/bounding_box/multi_box.py

@@ -0,0 +1,22 @@
+from cvdatasets.dataset.mixins.base import BaseMixin
+from cvdatasets.dataset.mixins.bounding_box.bbox import BoundingBox
+
+class MultiBoxMixin(BaseMixin):
+	_all_keys=[
+		"x", "x0", "x1",
+		"y", "y0", "y1",
+		"w", "h",
+	]
+
+	def multi_box(self, i, keys=["x0","x1","y0","y1"]):
+		assert keys is None or all([key in self._all_keys for key in keys]), \
+			f"unknown keys found: {keys}. Possible are: {self._all_keys}"
+
+		boxes = [BoundingBox.new(**box)
+			for box in self._get("multi_box", i)["objects"]
+		]
+
+		if keys is None:
+			return boxes
+
+		return [box.get(*keys) for box in boxes]