processors.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import numpy as np
  15. from ....utils import logging
  16. from ..common.vision import F
  17. class Crop:
  18. """Crop region from the image."""
  19. def __init__(self, crop_size, mode="C"):
  20. """
  21. Initialize the instance.
  22. Args:
  23. crop_size (list|tuple|int): Width and height of the region to crop.
  24. mode (str, optional): 'C' for cropping the center part and 'TL' for
  25. cropping the top left part. Default: 'C'.
  26. """
  27. super().__init__()
  28. if isinstance(crop_size, int):
  29. crop_size = [crop_size, crop_size]
  30. F.check_image_size(crop_size)
  31. self.crop_size = crop_size
  32. if mode not in ("C", "TL"):
  33. raise ValueError("Unsupported interpolation method")
  34. self.mode = mode
  35. def __call__(self, imgs):
  36. """apply"""
  37. return [self.crop(img) for img in imgs]
  38. def crop(self, img):
  39. h, w = img.shape[:2]
  40. cw, ch = self.crop_size
  41. if self.mode == "C":
  42. x1 = max(0, (w - cw) // 2)
  43. y1 = max(0, (h - ch) // 2)
  44. elif self.mode == "TL":
  45. x1, y1 = 0, 0
  46. x2 = min(w, x1 + cw)
  47. y2 = min(h, y1 + ch)
  48. coords = (x1, y1, x2, y2)
  49. if coords == (0, 0, w, h):
  50. raise ValueError(
  51. f"Input image ({w}, {h}) smaller than the target size ({cw}, {ch})."
  52. )
  53. img = F.slice(img, coords=coords)
  54. return img
  55. class Topk:
  56. """Topk Transform"""
  57. def __init__(self, topk, class_ids=None):
  58. super().__init__()
  59. assert isinstance(topk, (int,))
  60. self.topk = topk
  61. self.class_id_map = self._parse_class_id_map(class_ids)
  62. def _parse_class_id_map(self, class_ids):
  63. """parse class id to label map file"""
  64. if class_ids is None:
  65. return None
  66. class_id_map = {id: str(lb) for id, lb in enumerate(class_ids)}
  67. return class_id_map
  68. def __call__(self, preds):
  69. indexes = preds[0].argsort(axis=1)[:, -self.topk :][:, ::-1].astype("int32")
  70. scores = [
  71. np.around(pred[index], decimals=5) for pred, index in zip(preds[0], indexes)
  72. ]
  73. label_names = [[self.class_id_map[i] for i in index] for index in indexes]
  74. return indexes, scores, label_names