result.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  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 copy
  15. import cv2
  16. import numpy as np
  17. from PIL import Image
  18. from ...common.result import BaseCVResult, JsonMixin
  19. from ...utils.color_map import get_colormap
  20. from ..object_detection.result import draw_box
  21. def draw_segm(im, masks, mask_info, alpha=0.7):
  22. """
  23. Draw segmentation on image
  24. """
  25. w_ratio = 0.4
  26. color_list = get_colormap(rgb=True)
  27. im = np.array(im).astype("float32")
  28. clsid2color = {}
  29. masks = np.array(masks)
  30. masks = masks.astype(np.uint8)
  31. for i in range(masks.shape[0]):
  32. mask, score, clsid = masks[i], mask_info[i]["score"], mask_info[i]["class_id"]
  33. if clsid not in clsid2color:
  34. color_index = i % len(color_list)
  35. clsid2color[clsid] = color_list[color_index]
  36. color_mask = clsid2color[clsid]
  37. for c in range(3):
  38. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  39. idx = np.nonzero(mask)
  40. color_mask = np.array(color_mask)
  41. idx0 = np.minimum(idx[0], im.shape[0] - 1)
  42. idx1 = np.minimum(idx[1], im.shape[1] - 1)
  43. im[idx0, idx1, :] *= 1.0 - alpha
  44. im[idx0, idx1, :] += alpha * color_mask
  45. sum_x = np.sum(mask, axis=0)
  46. x = np.where(sum_x > 0.5)[0]
  47. sum_y = np.sum(mask, axis=1)
  48. y = np.where(sum_y > 0.5)[0]
  49. x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1]
  50. cv2.rectangle(
  51. im, (x0, y0), (x1, y1), tuple(color_mask.astype("int32").tolist()), 1
  52. )
  53. bbox_text = "%s %.2f" % (mask_info[i]["label"], score)
  54. t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0]
  55. cv2.rectangle(
  56. im,
  57. (x0, y0),
  58. (x0 + t_size[0], y0 - t_size[1] - 3),
  59. tuple(color_mask.astype("int32").tolist()),
  60. -1,
  61. )
  62. cv2.putText(
  63. im,
  64. bbox_text,
  65. (x0, y0 - 2),
  66. cv2.FONT_HERSHEY_SIMPLEX,
  67. 0.3,
  68. (0, 0, 0),
  69. 1,
  70. lineType=cv2.LINE_AA,
  71. )
  72. return Image.fromarray(im.astype("uint8"))
  73. def restore_to_draw_masks(img_size, boxes, masks):
  74. """
  75. Restores extracted masks to the original shape and draws them on a blank image.
  76. """
  77. restored_masks = []
  78. for i, (box, mask) in enumerate(zip(boxes, masks)):
  79. restored_mask = np.zeros(img_size, dtype=np.uint8)
  80. x_min, y_min, x_max, y_max = map(lambda x: int(round(x)), box["coordinate"])
  81. restored_mask[y_min:y_max, x_min:x_max] = mask
  82. restored_masks.append(restored_mask)
  83. return np.array(restored_masks)
  84. def draw_mask(im, boxes, np_masks, img_size):
  85. """
  86. Args:
  87. im (PIL.Image.Image): PIL image
  88. boxes (list): a list of dictionaries representing detection box information.
  89. np_masks (np.ndarray): shape:[N, im_h, im_w]
  90. Returns:
  91. im (PIL.Image.Image): visualized image
  92. """
  93. color_list = get_colormap(rgb=True)
  94. w_ratio = 0.4
  95. alpha = 0.7
  96. im = np.array(im).astype("float32")
  97. clsid2color = {}
  98. np_masks = restore_to_draw_masks(img_size, boxes, np_masks)
  99. im_h, im_w = im.shape[:2]
  100. np_masks = np_masks[:, :im_h, :im_w]
  101. for i in range(len(np_masks)):
  102. clsid, score = int(boxes[i]["cls_id"]), boxes[i]["score"]
  103. mask = np_masks[i]
  104. if clsid not in clsid2color:
  105. color_index = i % len(color_list)
  106. clsid2color[clsid] = color_list[color_index]
  107. color_mask = clsid2color[clsid]
  108. for c in range(3):
  109. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  110. idx = np.nonzero(mask)
  111. color_mask = np.array(color_mask)
  112. im[idx[0], idx[1], :] *= 1.0 - alpha
  113. im[idx[0], idx[1], :] += alpha * color_mask
  114. return Image.fromarray(im.astype("uint8"))
  115. class InstanceSegResult(BaseCVResult):
  116. """Save Result Transform"""
  117. def _to_img(self):
  118. """apply"""
  119. # image = self._img_reader.read(self["input_path"])
  120. image = Image.fromarray(self["input_img"])
  121. ori_img_size = list(image.size)[::-1]
  122. boxes = self["boxes"]
  123. masks = self["masks"]
  124. if next((True for item in self["boxes"] if "coordinate" in item), False):
  125. image = draw_mask(image, boxes, masks, ori_img_size)
  126. image = draw_box(image, boxes)
  127. else:
  128. image = draw_segm(image, masks, boxes)
  129. return {"res": image}
  130. def _to_str(self, *args, **kwargs):
  131. data = copy.deepcopy(self)
  132. data.pop("input_img")
  133. data["masks"] = "..."
  134. return JsonMixin._to_str(data, *args, **kwargs)
  135. def _to_json(self, *args, **kwargs):
  136. data = copy.deepcopy(self)
  137. data.pop("input_img")
  138. return JsonMixin._to_json(data, *args, **kwargs)