sam_result.py 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145
  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 random
  16. import cv2
  17. import numpy as np
  18. from PIL import Image
  19. from ....common.result import BaseCVResult, JsonMixin
  20. from ....utils.color_map import get_colormap
  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 = masks[i]
  33. clsid = random.randint(0, len(get_colormap(rgb=True)) - 1)
  34. if clsid not in clsid2color:
  35. color_index = i % len(color_list)
  36. clsid2color[clsid] = color_list[color_index]
  37. color_mask = clsid2color[clsid]
  38. for c in range(3):
  39. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  40. idx = np.nonzero(mask)
  41. color_mask = np.array(color_mask)
  42. idx0 = np.minimum(idx[0], im.shape[0] - 1)
  43. idx1 = np.minimum(idx[1], im.shape[1] - 1)
  44. im[idx0, idx1, :] *= 1.0 - alpha
  45. im[idx0, idx1, :] += alpha * color_mask
  46. # draw box prompt
  47. if mask_info[i]["label"] == "box_prompt":
  48. x0, y0, x1, y1 = mask_info[i]["prompt"]
  49. x0, y0, x1, y1 = int(x0), int(y0), int(x1), int(y1)
  50. cv2.rectangle(
  51. im, (x0, y0), (x1, y1), tuple(color_mask.astype("int32").tolist()), 1
  52. )
  53. bbox_text = "%s" % mask_info[i]["label"]
  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. elif mask_info[i]["label"] == "point_prompt":
  73. x, y = mask_info[i]["prompt"]
  74. bbox_text = "%s" % mask_info[i]["label"]
  75. t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0]
  76. cv2.circle(
  77. im,
  78. (x, y),
  79. 1,
  80. (255, 255, 255),
  81. 4,
  82. )
  83. cv2.putText(
  84. im,
  85. bbox_text,
  86. (x - t_size[0] // 2, y - t_size[1] - 1),
  87. cv2.FONT_HERSHEY_SIMPLEX,
  88. 0.3,
  89. (255, 255, 255),
  90. 1,
  91. lineType=cv2.LINE_AA,
  92. )
  93. else:
  94. raise NotImplementedError(
  95. f"Prompt type {mask_info[i]['label']} not implemented."
  96. )
  97. return Image.fromarray(im.astype("uint8"))
  98. class SAMSegResult(BaseCVResult):
  99. """Save Result Transform for SAM"""
  100. def __init__(self, data: dict) -> None:
  101. data["masks"] = [mask.squeeze(0) for mask in list(data["masks"])]
  102. prompts = data["prompts"]
  103. assert isinstance(prompts, dict) and len(prompts) == 1
  104. prompt_type, prompts = list(prompts.items())[0]
  105. mask_infos = [
  106. {
  107. "label": prompt_type,
  108. "prompt": p,
  109. }
  110. for p in prompts
  111. ]
  112. data["mask_infos"] = mask_infos
  113. assert len(data["masks"]) == len(mask_infos)
  114. super().__init__(data)
  115. def _to_img(self):
  116. """apply"""
  117. image = Image.fromarray(self["input_img"])
  118. mask_infos = self["mask_infos"]
  119. masks = self["masks"]
  120. image = draw_segm(image, masks, mask_infos)
  121. return {"res": image}
  122. def _to_str(self, *args, **kwargs):
  123. data = copy.deepcopy(self)
  124. data.pop("input_img")
  125. data["masks"] = "..."
  126. return JsonMixin._to_str(data, *args, **kwargs)
  127. def _to_json(self, *args, **kwargs):
  128. data = copy.deepcopy(self)
  129. data.pop("input_img")
  130. return JsonMixin._to_json(data, *args, **kwargs)