processors.py 7.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235
  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 os
  15. import ast
  16. import math
  17. from pathlib import Path
  18. from copy import deepcopy
  19. import numpy as np
  20. import cv2
  21. from . import funcs as F
  22. class _BaseResize:
  23. _INTERP_DICT = {
  24. "NEAREST": cv2.INTER_NEAREST,
  25. "LINEAR": cv2.INTER_LINEAR,
  26. "CUBIC": cv2.INTER_CUBIC,
  27. "AREA": cv2.INTER_AREA,
  28. "LANCZOS4": cv2.INTER_LANCZOS4,
  29. }
  30. def __init__(self, size_divisor, interp):
  31. super().__init__()
  32. if size_divisor is not None:
  33. assert isinstance(
  34. size_divisor, int
  35. ), "`size_divisor` should be None or int."
  36. self.size_divisor = size_divisor
  37. try:
  38. interp = self._INTERP_DICT[interp]
  39. except KeyError:
  40. raise ValueError(
  41. "`interp` should be one of {}.".format(self._INTERP_DICT.keys())
  42. )
  43. self.interp = interp
  44. @staticmethod
  45. def _rescale_size(img_size, target_size):
  46. """rescale size"""
  47. scale = min(max(target_size) / max(img_size), min(target_size) / min(img_size))
  48. rescaled_size = [round(i * scale) for i in img_size]
  49. return rescaled_size, scale
  50. class Resize(_BaseResize):
  51. """Resize the image."""
  52. def __init__(
  53. self, target_size, keep_ratio=False, size_divisor=None, interp="LINEAR"
  54. ):
  55. """
  56. Initialize the instance.
  57. Args:
  58. target_size (list|tuple|int): Target width and height.
  59. keep_ratio (bool, optional): Whether to keep the aspect ratio of resized
  60. image. Default: False.
  61. size_divisor (int|None, optional): Divisor of resized image size.
  62. Default: None.
  63. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  64. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  65. """
  66. super().__init__(size_divisor=size_divisor, interp=interp)
  67. if isinstance(target_size, int):
  68. target_size = [target_size, target_size]
  69. F.check_image_size(target_size)
  70. self.target_size = target_size
  71. self.keep_ratio = keep_ratio
  72. def __call__(self, imgs):
  73. """apply"""
  74. return [self.resize(img) for img in imgs]
  75. def resize(self, img):
  76. target_size = self.target_size
  77. original_size = img.shape[:2][::-1]
  78. if self.keep_ratio:
  79. h, w = img.shape[0:2]
  80. target_size, _ = self._rescale_size((w, h), self.target_size)
  81. if self.size_divisor:
  82. target_size = [
  83. math.ceil(i / self.size_divisor) * self.size_divisor
  84. for i in target_size
  85. ]
  86. img = F.resize(img, target_size, interp=self.interp)
  87. return img
  88. class ResizeByLong(_BaseResize):
  89. """
  90. Proportionally resize the image by specifying the target length of the
  91. longest side.
  92. """
  93. def __init__(self, target_long_edge, size_divisor=None, interp="LINEAR"):
  94. """
  95. Initialize the instance.
  96. Args:
  97. target_long_edge (int): Target length of the longest side of image.
  98. size_divisor (int|None, optional): Divisor of resized image size.
  99. Default: None.
  100. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  101. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  102. """
  103. super().__init__(size_divisor=size_divisor, interp=interp)
  104. self.target_long_edge = target_long_edge
  105. def __call__(self, imgs):
  106. """apply"""
  107. return [self.resize(img) for img in imgs]
  108. def resize(self, img):
  109. h, w = img.shape[:2]
  110. scale = self.target_long_edge / max(h, w)
  111. h_resize = round(h * scale)
  112. w_resize = round(w * scale)
  113. if self.size_divisor is not None:
  114. h_resize = math.ceil(h_resize / self.size_divisor) * self.size_divisor
  115. w_resize = math.ceil(w_resize / self.size_divisor) * self.size_divisor
  116. img = F.resize(img, (w_resize, h_resize), interp=self.interp)
  117. return img
  118. class ResizeByShort(_BaseResize):
  119. """
  120. Proportionally resize the image by specifying the target length of the
  121. shortest side.
  122. """
  123. def __init__(self, target_short_edge, size_divisor=None, interp="LINEAR"):
  124. """
  125. Initialize the instance.
  126. Args:
  127. target_short_edge (int): Target length of the shortest side of image.
  128. size_divisor (int|None, optional): Divisor of resized image size.
  129. Default: None.
  130. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  131. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  132. """
  133. super().__init__(size_divisor=size_divisor, interp=interp)
  134. self.target_short_edge = target_short_edge
  135. def __call__(self, imgs):
  136. """apply"""
  137. return [self.resize(img) for img in imgs]
  138. def resize(self, img):
  139. h, w = img.shape[:2]
  140. scale = self.target_short_edge / min(h, w)
  141. h_resize = round(h * scale)
  142. w_resize = round(w * scale)
  143. if self.size_divisor is not None:
  144. h_resize = math.ceil(h_resize / self.size_divisor) * self.size_divisor
  145. w_resize = math.ceil(w_resize / self.size_divisor) * self.size_divisor
  146. img = F.resize(img, (w_resize, h_resize), interp=self.interp)
  147. return img
  148. class Normalize:
  149. """Normalize the image."""
  150. def __init__(self, scale=1.0 / 255, mean=0.5, std=0.5, preserve_dtype=False):
  151. """
  152. Initialize the instance.
  153. Args:
  154. scale (float, optional): Scaling factor to apply to the image before
  155. applying normalization. Default: 1/255.
  156. mean (float|tuple|list, optional): Means for each channel of the image.
  157. Default: 0.5.
  158. std (float|tuple|list, optional): Standard deviations for each channel
  159. of the image. Default: 0.5.
  160. preserve_dtype (bool, optional): Whether to preserve the original dtype
  161. of the image.
  162. """
  163. super().__init__()
  164. self.scale = np.float32(scale)
  165. if isinstance(mean, float):
  166. mean = [mean]
  167. self.mean = np.asarray(mean).astype("float32")
  168. if isinstance(std, float):
  169. std = [std]
  170. self.std = np.asarray(std).astype("float32")
  171. self.preserve_dtype = preserve_dtype
  172. def __call__(self, imgs):
  173. """apply"""
  174. old_type = imgs[0].dtype
  175. # XXX: If `old_type` has higher precision than float32,
  176. # we will lose some precision.
  177. imgs = np.array(imgs).astype("float32", copy=False)
  178. imgs *= self.scale
  179. imgs -= self.mean
  180. imgs /= self.std
  181. if self.preserve_dtype:
  182. imgs = imgs.astype(old_type, copy=False)
  183. return list(imgs)
  184. class ToCHWImage:
  185. """Reorder the dimensions of the image from HWC to CHW."""
  186. def __call__(self, imgs):
  187. """apply"""
  188. return [img.transpose((2, 0, 1)) for img in imgs]
  189. class ToBatch:
  190. def __call__(self, imgs):
  191. return [np.stack(imgs, axis=0).astype(dtype=np.float32, copy=False)]