processors.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  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 ...utils.benchmark import benchmark
  18. from ..common.vision import funcs as F
  19. from ..common.vision.processors import _BaseResize
  20. @benchmark.timeit
  21. class Resize(_BaseResize):
  22. """Resize the image."""
  23. def __init__(
  24. self, target_size=-1, keep_ratio=False, size_divisor=None, interp="LINEAR"
  25. ):
  26. """
  27. Initialize the instance.
  28. Args:
  29. target_size (list|tuple|int, optional): Target width and height. -1 will return the images directly without resizing.
  30. keep_ratio (bool, optional): Whether to keep the aspect ratio of resized
  31. image. Default: False.
  32. size_divisor (int|None, optional): Divisor of resized image size.
  33. Default: None.
  34. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  35. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  36. """
  37. super().__init__(size_divisor=size_divisor, interp=interp)
  38. if isinstance(target_size, int):
  39. target_size = (target_size, target_size)
  40. F.check_image_size(target_size)
  41. self.target_size = target_size
  42. self.keep_ratio = keep_ratio
  43. def __call__(self, imgs, target_size=None):
  44. """apply"""
  45. target_size = self.target_size if target_size is None else target_size
  46. if isinstance(target_size, int):
  47. target_size = (target_size, target_size)
  48. F.check_image_size(target_size)
  49. return [self.resize(img, target_size) for img in imgs]
  50. def resize(self, img, target_size):
  51. if target_size == (-1, -1):
  52. # If the final target_size == (-1, -1), it means use the source input image directly.
  53. return img
  54. original_size = img.shape[:2][::-1]
  55. assert target_size[0] > 0 and target_size[1] > 0
  56. if self.keep_ratio:
  57. h, w = img.shape[0:2]
  58. target_size, _ = self._rescale_size((w, h), target_size)
  59. if self.size_divisor:
  60. target_size = [
  61. math.ceil(i / self.size_divisor) * self.size_divisor
  62. for i in target_size
  63. ]
  64. img = F.resize(img, target_size, interp=self.interp)
  65. return img
  66. @benchmark.timeit
  67. class SegPostProcess:
  68. """Semantic Segmentation PostProcess
  69. This class is responsible for post-processing detection results, only including
  70. restoring the prediction segmentation map to the original image size for now.
  71. """
  72. def __call__(self, imgs, src_images):
  73. assert len(imgs) == len(src_images)
  74. src_sizes = [src_image.shape[:2][::-1] for src_image in src_images]
  75. return [
  76. self.reverse_resize(img, src_size) for img, src_size in zip(imgs, src_sizes)
  77. ]
  78. def reverse_resize(self, img, src_size):
  79. """Restore the prediction map to source image size using nearest interpolation.
  80. Args:
  81. img (np.ndarray): prediction map with shape of (1, width, height)
  82. src_size (Tuple[int, int]): source size of the input image, with format of (width, height).
  83. """
  84. assert isinstance(src_size, (tuple, list)) and len(src_size) == 2
  85. assert src_size[0] > 0 and src_size[1] > 0
  86. assert img.ndim == 3
  87. reversed_img = cv2.resize(
  88. img[0], dsize=src_size, interpolation=cv2.INTER_NEAREST
  89. )
  90. reversed_img = np.expand_dims(reversed_img, axis=0)
  91. return reversed_img