common.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109
  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. from typing import Dict, List, Tuple
  15. import cv2
  16. import numpy as np
  17. class LetterResize(object):
  18. def __init__(
  19. self,
  20. scale=[640, 640],
  21. pad_val=144,
  22. use_mini_pad=False,
  23. stretch_only=False,
  24. allow_scale_up=False,
  25. ):
  26. super(LetterResize, self).__init__()
  27. self.scale = scale
  28. self.pad_val = pad_val
  29. self.use_mini_pad = use_mini_pad
  30. self.stretch_only = stretch_only
  31. self.allow_scale_up = allow_scale_up
  32. def _resize_img(self, image: np.ndarray) -> Dict:
  33. scale = self.scale
  34. image_shape = image.shape[:2]
  35. ratio = min(scale[0] / image_shape[0], scale[1] / image_shape[1])
  36. if not self.allow_scale_up:
  37. ratio = min(ratio, 1.0)
  38. ratio = [ratio, ratio]
  39. no_pad_shape = (
  40. int(round(image_shape[0] * ratio[0])),
  41. int(round(image_shape[1] * ratio[1])),
  42. )
  43. padding_h, padding_w = [scale[0] - no_pad_shape[0], scale[1] - no_pad_shape[1]]
  44. if self.use_mini_pad:
  45. padding_w, padding_h = np.mod(padding_w, 32), np.mod(padding_h, 32)
  46. elif self.stretch_only:
  47. padding_h, padding_w = 0.0, 0.0
  48. no_pad_shape = (scale[0], scale[1])
  49. ratio = [scale[0] / image_shape[0], scale[1] / image_shape[1]]
  50. if image_shape != no_pad_shape:
  51. image = cv2.resize(
  52. image,
  53. (no_pad_shape[1], no_pad_shape[0]),
  54. interpolation=cv2.INTER_LINEAR,
  55. )
  56. scale_factor = (
  57. no_pad_shape[1] / image_shape[1],
  58. no_pad_shape[0] / image_shape[0],
  59. )
  60. top_padding, left_padding = int(round(padding_h // 2 - 0.1)), int(
  61. round(padding_w // 2 - 0.1)
  62. )
  63. bottom_padding = padding_h - top_padding
  64. right_padding = padding_w - left_padding
  65. padding_list = [top_padding, bottom_padding, left_padding, right_padding]
  66. if (
  67. top_padding != 0
  68. or bottom_padding != 0
  69. or left_padding != 0
  70. or right_padding != 0
  71. ):
  72. pad_val = self.pad_val
  73. if isinstance(pad_val, int) and image.ndim == 3:
  74. pad_val = tuple(pad_val for _ in range(image.shape[2]))
  75. top, bottom, left, right = padding_list
  76. image = cv2.copyMakeBorder(
  77. image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=pad_val
  78. )
  79. result = dict()
  80. result["image"] = image
  81. result["scale_factor"] = np.array(scale_factor, dtype=np.float32)
  82. result["pad_param"] = np.array(padding_list, dtype=np.float32)
  83. return result
  84. def __call__(self, images: List[np.ndarray]) -> List[Dict]:
  85. if not isinstance(images, (List, Tuple)):
  86. images = [images]
  87. rst_images = [self._resize_img(image) for image in images]
  88. return rst_images