collection.py 969 B

12345678910111213141516171819202122232425262728293031323334353637
  1. import numpy as np
  2. from nabirds import utils
  3. from .base import BasePartCollection, BasePart
  4. from .annotation import BBoxPart
  5. class Parts(BasePartCollection):
  6. def __init__(self, image, part_annotations, rescale_size):
  7. super(Parts, self).__init__()
  8. self._parts = [BasePart.new(image, a, rescale_size) for a in part_annotations]
  9. class UniformParts(BasePartCollection):
  10. def __init__(self, image, ratio):
  11. super(UniformParts, self).__init__()
  12. self._parts = list(self.generate_parts(image, ratio))
  13. def generate_parts(self, im, ratio, round_op=np.floor):
  14. h, w, c = utils.dimensions(im)
  15. part_w = round_op(w * ratio).astype(np.int32)
  16. part_h = round_op(h * ratio).astype(np.int32)
  17. n, m = w // part_w, h // part_h
  18. # fit best possible part_w and part_h
  19. part_w = int(w / n)
  20. part_h = int(h / m)
  21. for i in range(n*m):
  22. row, col = np.unravel_index(i, (n, m))
  23. x, y = col * part_w, row * part_h
  24. yield BBoxPart(im, [i, x, y, part_w, part_h])