transforms.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. import os
  12. import numpy as np
  13. import math
  14. from PIL import Image, ImageDraw, ImageFont
  15. from .keys import DetKeys as K
  16. from ...base import BaseTransform
  17. from ...base.predictor.io.writers import ImageWriter
  18. from ...base.predictor.transforms import image_functions as F
  19. from ...base.predictor.transforms.image_common import _BaseResize, _check_image_size
  20. from ....utils.fonts import PINGFANG_FONT_FILE_PATH
  21. from ....utils import logging
  22. __all__ = ['SaveDetResults', 'PadStride', 'DetResize', 'PrintResult']
  23. def get_color_map_list(num_classes):
  24. """
  25. Args:
  26. num_classes (int): number of class
  27. Returns:
  28. color_map (list): RGB color list
  29. """
  30. color_map = num_classes * [0, 0, 0]
  31. for i in range(0, num_classes):
  32. j = 0
  33. lab = i
  34. while lab:
  35. color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
  36. color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
  37. color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
  38. j += 1
  39. lab >>= 3
  40. color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
  41. return color_map
  42. def colormap(rgb=False):
  43. """
  44. Get colormap
  45. The code of this function is copied from https://github.com/facebookresearch/Detectron/blob/main/detectron/\
  46. utils/colormap.py
  47. """
  48. color_list = np.array([
  49. 0xFF, 0x00, 0x00, 0xCC, 0xFF, 0x00, 0x00, 0xFF, 0x66, 0x00, 0x66, 0xFF,
  50. 0xCC, 0x00, 0xFF, 0xFF, 0x4D, 0x00, 0x80, 0xff, 0x00, 0x00, 0xFF, 0xB2,
  51. 0x00, 0x1A, 0xFF, 0xFF, 0x00, 0xE5, 0xFF, 0x99, 0x00, 0x33, 0xFF, 0x00,
  52. 0x00, 0xFF, 0xFF, 0x33, 0x00, 0xFF, 0xff, 0x00, 0x99, 0xFF, 0xE5, 0x00,
  53. 0x00, 0xFF, 0x1A, 0x00, 0xB2, 0xFF, 0x80, 0x00, 0xFF, 0xFF, 0x00, 0x4D
  54. ]).astype(np.float32)
  55. color_list = (color_list.reshape((-1, 3)))
  56. if not rgb:
  57. color_list = color_list[:, ::-1]
  58. return color_list.astype('int32')
  59. def font_colormap(color_index):
  60. """
  61. Get font color according to the index of colormap
  62. """
  63. dark = np.array([0x14, 0x0E, 0x35])
  64. light = np.array([0xFF, 0xFF, 0xFF])
  65. light_indexs = [0, 3, 4, 8, 9, 13, 14, 18, 19]
  66. if color_index in light_indexs:
  67. return light.astype('int32')
  68. else:
  69. return dark.astype('int32')
  70. def draw_box(img, np_boxes, labels, threshold=0.5):
  71. """
  72. Args:
  73. img (PIL.Image.Image): PIL image
  74. np_boxes (np.ndarray): shape:[N,6], N: number of box,
  75. matix element:[class, score, x_min, y_min, x_max, y_max]
  76. labels (list): labels:['class1', ..., 'classn']
  77. threshold (float): threshold of box
  78. Returns:
  79. img (PIL.Image.Image): visualized image
  80. """
  81. font_size = int(0.024 * int(img.width)) + 2
  82. font = ImageFont.truetype(
  83. PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
  84. draw_thickness = int(max(img.size) * 0.005)
  85. draw = ImageDraw.Draw(img)
  86. clsid2color = {}
  87. catid2fontcolor = {}
  88. color_list = colormap(rgb=True)
  89. expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
  90. np_boxes = np_boxes[expect_boxes, :]
  91. for i, dt in enumerate(np_boxes):
  92. clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
  93. if clsid not in clsid2color:
  94. color_index = i % len(color_list)
  95. clsid2color[clsid] = color_list[color_index]
  96. catid2fontcolor[clsid] = font_colormap(color_index)
  97. color = tuple(clsid2color[clsid])
  98. font_color = tuple(catid2fontcolor[clsid])
  99. xmin, ymin, xmax, ymax = bbox
  100. # draw bbox
  101. draw.line(
  102. [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
  103. (xmin, ymin)],
  104. width=draw_thickness,
  105. fill=color)
  106. # draw label
  107. text = "{} {:.2f}".format(labels[clsid], score)
  108. tw, th = draw.textsize(text, font=font)
  109. if ymin < th:
  110. draw.rectangle(
  111. [(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
  112. draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
  113. else:
  114. draw.rectangle(
  115. [(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
  116. draw.text(
  117. (xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
  118. return img
  119. def draw_mask(im, np_boxes, np_masks, labels, threshold=0.5):
  120. """
  121. Args:
  122. im (PIL.Image.Image): PIL image
  123. np_boxes (np.ndarray): shape:[N,6], N: number of box,
  124. matix element:[class, score, x_min, y_min, x_max, y_max]
  125. np_masks (np.ndarray): shape:[N, im_h, im_w]
  126. labels (list): labels:['class1', ..., 'classn']
  127. threshold (float): threshold of mask
  128. Returns:
  129. im (PIL.Image.Image): visualized image
  130. """
  131. color_list = get_color_map_list(len(labels))
  132. w_ratio = 0.4
  133. alpha = 0.7
  134. im = np.array(im).astype('float32')
  135. clsid2color = {}
  136. expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
  137. np_boxes = np_boxes[expect_boxes, :]
  138. np_masks = np_masks[expect_boxes, :, :]
  139. im_h, im_w = im.shape[:2]
  140. np_masks = np_masks[:, :im_h, :im_w]
  141. for i in range(len(np_masks)):
  142. clsid, score = int(np_boxes[i][0]), np_boxes[i][1]
  143. mask = np_masks[i]
  144. if clsid not in clsid2color:
  145. clsid2color[clsid] = color_list[clsid]
  146. color_mask = clsid2color[clsid]
  147. for c in range(3):
  148. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  149. idx = np.nonzero(mask)
  150. color_mask = np.array(color_mask)
  151. im[idx[0], idx[1], :] *= 1.0 - alpha
  152. im[idx[0], idx[1], :] += alpha * color_mask
  153. return Image.fromarray(im.astype('uint8'))
  154. class SaveDetResults(BaseTransform):
  155. """ Save Result Transform """
  156. def __init__(self, save_dir, threshold=0.5, labels=None):
  157. super().__init__()
  158. self.save_dir = save_dir
  159. self.threshold = threshold
  160. self.labels = labels
  161. # We use pillow backend to save both numpy arrays and PIL Image objects
  162. self._writer = ImageWriter(backend='pillow')
  163. def apply(self, data):
  164. """ apply """
  165. ori_path = data[K.IM_PATH]
  166. file_name = os.path.basename(ori_path)
  167. save_path = os.path.join(self.save_dir, file_name)
  168. labels = self.labels
  169. image = Image.open(ori_path)
  170. if K.MASKS in data:
  171. image = draw_mask(
  172. image,
  173. data[K.BOXES],
  174. data[K.MASKS],
  175. threshold=self.threshold,
  176. labels=labels)
  177. image = draw_box(
  178. image, data[K.BOXES], threshold=self.threshold, labels=labels)
  179. self._write_image(save_path, image)
  180. return data
  181. def _write_image(self, path, image):
  182. """ write image """
  183. if os.path.exists(path):
  184. logging.warning(f"{path} already exists. Overwriting it.")
  185. self._writer.write(path, image)
  186. @classmethod
  187. def get_input_keys(cls):
  188. """ get input keys """
  189. return [K.IM_PATH, K.BOXES]
  190. @classmethod
  191. def get_output_keys(cls):
  192. """ get output keys """
  193. return []
  194. class PadStride(BaseTransform):
  195. """ padding image for model with FPN , instead PadBatch(pad_to_stride, pad_gt) in original config
  196. Args:
  197. stride (bool): model with FPN need image shape % stride == 0
  198. """
  199. def __init__(self, stride=0):
  200. self.coarsest_stride = stride
  201. def apply(self, data):
  202. """
  203. Args:
  204. im (np.ndarray): image (np.ndarray)
  205. Returns:
  206. im (np.ndarray): processed image (np.ndarray)
  207. """
  208. im = data[K.IMAGE]
  209. coarsest_stride = self.coarsest_stride
  210. if coarsest_stride <= 0:
  211. return data
  212. im_c, im_h, im_w = im.shape
  213. pad_h = int(np.ceil(float(im_h) / coarsest_stride) * coarsest_stride)
  214. pad_w = int(np.ceil(float(im_w) / coarsest_stride) * coarsest_stride)
  215. padding_im = np.zeros((im_c, pad_h, pad_w), dtype=np.float32)
  216. padding_im[:, :im_h, :im_w] = im
  217. data[K.IMAGE] = padding_im
  218. return data
  219. @classmethod
  220. def get_input_keys(cls):
  221. """ get input keys """
  222. return [K.IMAGE]
  223. @classmethod
  224. def get_output_keys(cls):
  225. """ get output keys """
  226. return [K.IMAGE]
  227. class DetResize(_BaseResize):
  228. """
  229. Resize the image.
  230. Args:
  231. target_size (list|tuple|int): Target height and width.
  232. keep_ratio (bool, optional): Whether to keep the aspect ratio of resized
  233. image. Default: False.
  234. size_divisor (int|None, optional): Divisor of resized image size.
  235. Default: None.
  236. interp (str, optional): Interpolation method. Choices are 'NEAREST',
  237. 'LINEAR', 'CUBIC', 'AREA', and 'LANCZOS4'. Default: 'LINEAR'.
  238. """
  239. def __init__(self,
  240. target_hw,
  241. keep_ratio=False,
  242. size_divisor=None,
  243. interp='LINEAR'):
  244. super().__init__(size_divisor=size_divisor, interp=interp)
  245. if isinstance(target_hw, int):
  246. target_hw = [target_hw, target_hw]
  247. _check_image_size(target_hw)
  248. self.target_hw = target_hw
  249. self.keep_ratio = keep_ratio
  250. def apply(self, data):
  251. """ apply """
  252. target_hw = self.target_hw
  253. im = data['image']
  254. original_size = im.shape[:2]
  255. if self.keep_ratio:
  256. h, w = im.shape[0:2]
  257. target_hw, _ = self._rescale_size((h, w), self.target_hw)
  258. if self.size_divisor:
  259. target_hw = [
  260. math.ceil(i / self.size_divisor) * self.size_divisor
  261. for i in target_hw
  262. ]
  263. im_scale_w, im_scale_h = [
  264. target_hw[1] / original_size[1], target_hw[0] / original_size[0]
  265. ]
  266. im = F.resize(im, target_hw[::-1], interp=self.interp)
  267. data['image'] = im
  268. data['image_size'] = [im.shape[1], im.shape[0]]
  269. data['scale_factors'] = [im_scale_w, im_scale_h]
  270. return data
  271. @classmethod
  272. def get_input_keys(cls):
  273. """ get input keys """
  274. # image: Image in hw or hwc format.
  275. return ['image']
  276. @classmethod
  277. def get_output_keys(cls):
  278. """ get output keys """
  279. # image: Image in hw or hwc format.
  280. # image_size: Width and height of the image.
  281. # scale_factors: Scale factors for image width and height.
  282. return ['image', 'image_size', 'scale_factors']
  283. class PrintResult(BaseTransform):
  284. """ Print Result Transform """
  285. def apply(self, data):
  286. """ apply """
  287. logging.info("The prediction result is:")
  288. logging.info(data[K.BOXES])
  289. return data
  290. @classmethod
  291. def get_input_keys(cls):
  292. """ get input keys """
  293. return [K.BOXES]
  294. @classmethod
  295. def get_output_keys(cls):
  296. """ get output keys """
  297. return []