bbox_mixin.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. import abc
  2. import logging
  3. import numpy as np
  4. class BBoxMixin(abc.ABC):
  5. dtype = np.dtype([
  6. ("x", np.int32),
  7. ("y", np.int32),
  8. ("w", np.int32),
  9. ("h", np.int32),
  10. ])
  11. def read_annotation_files(self):
  12. files = super(BBoxMixin, self).read_annotation_files()
  13. files.load_files(
  14. bounding_boxes=("bounding_boxes.txt", True),
  15. )
  16. return files
  17. @property
  18. def has_bounding_boxes(self):
  19. return self.files.bounding_boxes is not None
  20. def parse_annotations(self):
  21. super(BBoxMixin, self).parse_annotations()
  22. if self.has_bounding_boxes:
  23. self._parse_bounding_boxes()
  24. def _parse_bounding_boxes(self):
  25. logging.debug("Parsing bounding box annotations")
  26. assert self.has_bounding_boxes, \
  27. "Bouding boxes were not loaded!"
  28. uuid_to_bbox = {}
  29. for content in [i.split() for i in self.files.bounding_boxes]:
  30. uuid, bbox = content[0], content[1:]
  31. uuid_to_bbox[uuid] = [float(i) for i in bbox]
  32. self.bounding_boxes = np.array(
  33. [tuple(uuid_to_bbox[uuid]) for uuid in self.uuids],
  34. dtype=self.dtype)
  35. def bounding_box(self, uuid):
  36. if self.has_bounding_boxes:
  37. return self.bounding_boxes[self.uuid_to_idx[uuid]].copy()
  38. return np.array((0,0, 1,1), dtype=self.dtype)