processors.py 8.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  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 math
  15. import cv2
  16. import numpy as np
  17. from PIL import Image
  18. from ....utils.benchmark import benchmark
  19. from . import funcs as F
  20. class _BaseResize:
  21. _CV2_INTERP_DICT = {
  22. "NEAREST": cv2.INTER_NEAREST,
  23. "LINEAR": cv2.INTER_LINEAR,
  24. "BICUBIC": cv2.INTER_CUBIC,
  25. "AREA": cv2.INTER_AREA,
  26. "LANCZOS4": cv2.INTER_LANCZOS4,
  27. }
  28. _PIL_INTERP_DICT = {
  29. "NEAREST": Image.NEAREST,
  30. "BILINEAR": Image.BILINEAR,
  31. "BICUBIC": Image.BICUBIC,
  32. "BOX": Image.BOX,
  33. "LANCZOS4": Image.LANCZOS,
  34. }
  35. def __init__(self, size_divisor, interp, backend="cv2"):
  36. super().__init__()
  37. if size_divisor is not None:
  38. assert isinstance(
  39. size_divisor, int
  40. ), "`size_divisor` should be None or int."
  41. self.size_divisor = size_divisor
  42. try:
  43. interp = interp.upper()
  44. if backend == "cv2":
  45. interp = self._CV2_INTERP_DICT[interp]
  46. elif backend == "pil":
  47. interp = self._PIL_INTERP_DICT[interp]
  48. else:
  49. raise ValueError("backend must be `cv2` or `pil`")
  50. except KeyError:
  51. raise ValueError(
  52. "For backend '{}', `interp` should be one of {}. Please ensure the interpolation method matches the selected backend.".format(
  53. backend,
  54. (
  55. self._CV2_INTERP_DICT.keys()
  56. if backend == "cv2"
  57. else self._PIL_INTERP_DICT.keys()
  58. ),
  59. )
  60. )
  61. self.interp = interp
  62. self.backend = backend
  63. @staticmethod
  64. def _rescale_size(img_size, target_size):
  65. """rescale size"""
  66. scale = min(max(target_size) / max(img_size), min(target_size) / min(img_size))
  67. rescaled_size = [round(i * scale) for i in img_size]
  68. return rescaled_size, scale
  69. @benchmark.timeit
  70. class Resize(_BaseResize):
  71. """Resize the image."""
  72. def __init__(
  73. self,
  74. target_size,
  75. keep_ratio=False,
  76. size_divisor=None,
  77. interp="LINEAR",
  78. backend="cv2",
  79. ):
  80. """
  81. Initialize the instance.
  82. Args:
  83. target_size (list|tuple|int): Target width and height.
  84. keep_ratio (bool, optional): Whether to keep the aspect ratio of resized
  85. image. Default: False.
  86. size_divisor (int|None, optional): Divisor of resized image size.
  87. Default: None.
  88. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  89. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  90. """
  91. super().__init__(size_divisor=size_divisor, interp=interp, backend=backend)
  92. if isinstance(target_size, int):
  93. target_size = [target_size, target_size]
  94. F.check_image_size(target_size)
  95. self.target_size = target_size
  96. self.keep_ratio = keep_ratio
  97. def __call__(self, imgs):
  98. """apply"""
  99. return [self.resize(img) for img in imgs]
  100. def resize(self, img):
  101. target_size = self.target_size
  102. original_size = img.shape[:2][::-1]
  103. if self.keep_ratio:
  104. h, w = img.shape[0:2]
  105. target_size, _ = self._rescale_size((w, h), self.target_size)
  106. if self.size_divisor:
  107. target_size = [
  108. math.ceil(i / self.size_divisor) * self.size_divisor
  109. for i in target_size
  110. ]
  111. img = F.resize(img, target_size, interp=self.interp, backend=self.backend)
  112. return img
  113. @benchmark.timeit
  114. class ResizeByLong(_BaseResize):
  115. """
  116. Proportionally resize the image by specifying the target length of the
  117. longest side.
  118. """
  119. def __init__(
  120. self, target_long_edge, size_divisor=None, interp="LINEAR", backend="cv2"
  121. ):
  122. """
  123. Initialize the instance.
  124. Args:
  125. target_long_edge (int): Target length of the longest side of image.
  126. size_divisor (int|None, optional): Divisor of resized image size.
  127. Default: None.
  128. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  129. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  130. """
  131. super().__init__(size_divisor=size_divisor, interp=interp, backend=backend)
  132. self.target_long_edge = target_long_edge
  133. def __call__(self, imgs):
  134. """apply"""
  135. return [self.resize(img) for img in imgs]
  136. def resize(self, img):
  137. h, w = img.shape[:2]
  138. scale = self.target_long_edge / max(h, w)
  139. h_resize = round(h * scale)
  140. w_resize = round(w * scale)
  141. if self.size_divisor is not None:
  142. h_resize = math.ceil(h_resize / self.size_divisor) * self.size_divisor
  143. w_resize = math.ceil(w_resize / self.size_divisor) * self.size_divisor
  144. img = F.resize(
  145. img, (w_resize, h_resize), interp=self.interp, backend=self.backend
  146. )
  147. return img
  148. @benchmark.timeit
  149. class ResizeByShort(_BaseResize):
  150. """
  151. Proportionally resize the image by specifying the target length of the
  152. shortest side.
  153. """
  154. def __init__(
  155. self, target_short_edge, size_divisor=None, interp="LINEAR", backend="cv2"
  156. ):
  157. """
  158. Initialize the instance.
  159. Args:
  160. target_short_edge (int): Target length of the shortest side of image.
  161. size_divisor (int|None, optional): Divisor of resized image size.
  162. Default: None.
  163. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  164. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  165. """
  166. super().__init__(size_divisor=size_divisor, interp=interp, backend=backend)
  167. self.target_short_edge = target_short_edge
  168. def __call__(self, imgs):
  169. """apply"""
  170. return [self.resize(img) for img in imgs]
  171. def resize(self, img):
  172. h, w = img.shape[:2]
  173. scale = self.target_short_edge / min(h, w)
  174. h_resize = round(h * scale)
  175. w_resize = round(w * scale)
  176. if self.size_divisor is not None:
  177. h_resize = math.ceil(h_resize / self.size_divisor) * self.size_divisor
  178. w_resize = math.ceil(w_resize / self.size_divisor) * self.size_divisor
  179. img = F.resize(
  180. img, (w_resize, h_resize), interp=self.interp, backend=self.backend
  181. )
  182. return img
  183. @benchmark.timeit
  184. class Normalize:
  185. """Normalize the three-channel image."""
  186. def __init__(self, scale=1.0 / 255, mean=0.5, std=0.5):
  187. """
  188. Initialize the instance.
  189. Args:
  190. scale (float, optional): Scaling factor to apply to the image before
  191. applying normalization. Default: 1/255.
  192. mean (float|tuple|list, optional): Means for each channel of the image.
  193. Default: 0.5.
  194. std (float|tuple|list|np.ndarray, optional): Standard deviations for each channel
  195. of the image. Default: 0.5.
  196. """
  197. super().__init__()
  198. if isinstance(mean, float):
  199. mean = [mean] * 3
  200. elif len(mean) != 3:
  201. raise ValueError(
  202. f"Expected `mean` to be a tuple or list of length 3, but got {len(mean)} elements."
  203. )
  204. if isinstance(std, float):
  205. std = [std] * 3
  206. elif len(std) != 3:
  207. raise ValueError(
  208. f"Expected `std` to be a tuple or list of length 3, but got {len(std)} elements."
  209. )
  210. self.alpha = [scale / std[i] for i in range(len(std))]
  211. self.beta = [-mean[i] / std[i] for i in range(len(std))]
  212. def norm(self, img):
  213. split_im = list(cv2.split(img))
  214. for c in range(img.shape[2]):
  215. split_im[c] = split_im[c].astype(np.float32)
  216. split_im[c] *= self.alpha[c]
  217. split_im[c] += self.beta[c]
  218. res = cv2.merge(split_im)
  219. return res
  220. def __call__(self, imgs):
  221. """apply"""
  222. return [self.norm(img) for img in imgs]
  223. @benchmark.timeit
  224. class ToCHWImage:
  225. """Reorder the dimensions of the image from HWC to CHW."""
  226. def __call__(self, imgs):
  227. """apply"""
  228. return [img.transpose((2, 0, 1)) for img in imgs]
  229. @benchmark.timeit
  230. class ToBatch:
  231. def __call__(self, imgs):
  232. return [np.stack(imgs, axis=0).astype(dtype=np.float32, copy=False)]