Browse Source

some import and indexing fixes

Dimitri Korsch 5 years ago
parent
commit
48c9189071

+ 1 - 1
cvdatasets/annotation/mixins/multi_box_mixin.py

@@ -43,7 +43,7 @@ class MultiBoxMixin(abc.ABC):
 
 	def multi_box(self, uuid) -> np.ndarray:
 		if self.has_multi_boxes:
-			return self.multi_boxes[self.uuid_to_idx[uuid]]
+			return self.multi_boxes[uuid]
 
 		fname = self.image_names[self.uuid_to_idx[uuid]]
 		return dict(image=fname, objects=[dict(x0=0, x1=0, y0=1, y1=1)])

+ 1 - 0
cvdatasets/dataset/__init__.py

@@ -3,6 +3,7 @@ from cvdatasets.dataset.mixins.features import PreExtractedFeaturesMixin
 from cvdatasets.dataset.mixins.parts import BBCropMixin
 from cvdatasets.dataset.mixins.parts import BBoxMixin
 from cvdatasets.dataset.mixins.parts import CroppedPartMixin
+from cvdatasets.dataset.mixins.parts import MultiBoxMixin
 from cvdatasets.dataset.mixins.parts import PartCropMixin
 from cvdatasets.dataset.mixins.parts import PartMixin
 from cvdatasets.dataset.mixins.parts import PartRevealMixin

+ 25 - 0
cvdatasets/dataset/mixins/parts.py

@@ -8,6 +8,31 @@ class BBoxMixin(BaseMixin):
 		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):