__init__.py 820 B

123456789101112131415161718192021222324252627282930313233
  1. from abc import ABC, abstractmethod
  2. import numpy as np
  3. import six
  4. from matplotlib.patches import Rectangle
  5. class BaseMixin(ABC):
  6. @abstractmethod
  7. def get_example(self, i):
  8. s = super(BaseMixin, self)
  9. if hasattr(s, "get_example"):
  10. return s.get_example(i)
  11. def plot_bounding_box(self, i, ax, fill=False, linestyle="--", **kwargs):
  12. x, y, w, h = self.bounding_box(i)
  13. ax.add_patch(Rectangle(
  14. (x,y), w, h,
  15. fill=False,
  16. linestyle="-.",
  17. **kwargs
  18. ))
  19. def __getitem__(self, index):
  20. if isinstance(index, slice):
  21. current, stop, step = index.indices(len(self))
  22. return [self.get_example(i) for i in
  23. six.moves.range(current, stop, step)]
  24. elif isinstance(index, list) or isinstance(index, np.ndarray):
  25. return [self.get_example(i) for i in index]
  26. else:
  27. return self.get_example(index)