12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849 |
- import abc
- import logging
- import numpy as np
- class BBoxMixin(abc.ABC):
- dtype = np.dtype([
- ("x", np.int32),
- ("y", np.int32),
- ("w", np.int32),
- ("h", np.int32),
- ])
- def read_annotation_files(self):
- files = super(BBoxMixin, self).read_annotation_files()
- files.load_files(
- bounding_boxes=("bounding_boxes.txt", True),
- )
- return files
- @property
- def has_bounding_boxes(self):
- return self.files.bounding_boxes is not None
- def parse_annotations(self):
- super(BBoxMixin, self).parse_annotations()
- if self.has_bounding_boxes:
- self._parse_bounding_boxes()
- def _parse_bounding_boxes(self):
- logging.debug("Parsing bounding box annotations")
- assert self.has_bounding_boxes, \
- "Bouding boxes were not loaded!"
- uuid_to_bbox = {}
- for content in [i.split() for i in self.files.bounding_boxes]:
- uuid, bbox = content[0], content[1:]
- uuid_to_bbox[uuid] = [float(i) for i in bbox]
- self.bounding_boxes = np.array(
- [tuple(uuid_to_bbox[uuid]) for uuid in self.uuids],
- dtype=self.dtype)
- def bounding_box(self, uuid):
- if self.has_bounding_boxes:
- return self.bounding_boxes[self.uuid_to_idx[uuid]].copy()
- return np.array((0,0, 1,1), dtype=self.dtype)
|