instance_seg.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  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 os
  15. import numpy as np
  16. import math
  17. import PIL
  18. from PIL import Image, ImageDraw, ImageFont
  19. from ...utils import logging
  20. from ...utils.fonts import PINGFANG_FONT_FILE_PATH
  21. from ..utils.io import ImageWriter, ImageReader
  22. from ..utils.color_map import get_colormap, font_colormap
  23. from .base import BaseResult
  24. from .det import draw_box
  25. def restore_to_draw_masks(img_size, boxes, masks):
  26. """
  27. Restores extracted masks to the original shape and draws them on a blank image.
  28. """
  29. restored_masks = []
  30. for i, (box, mask) in enumerate(zip(boxes, masks)):
  31. restored_mask = np.zeros(img_size, dtype=np.uint8)
  32. x_min, y_min, x_max, y_max = map(lambda x: int(round(x)), box["coordinate"])
  33. restored_mask[y_min:y_max, x_min:x_max] = mask
  34. restored_masks.append(restored_mask)
  35. return np.array(restored_masks)
  36. def draw_mask(im, boxes, np_masks, img_size):
  37. """
  38. Args:
  39. im (PIL.Image.Image): PIL image
  40. boxes (list): a list of dictionaries representing detection box information.
  41. np_masks (np.ndarray): shape:[N, im_h, im_w]
  42. Returns:
  43. im (PIL.Image.Image): visualized image
  44. """
  45. color_list = get_colormap(rgb=True)
  46. w_ratio = 0.4
  47. alpha = 0.7
  48. im = np.array(im).astype("float32")
  49. clsid2color = {}
  50. np_masks = restore_to_draw_masks(img_size, boxes, np_masks)
  51. im_h, im_w = im.shape[:2]
  52. np_masks = np_masks[:, :im_h, :im_w]
  53. for i in range(len(np_masks)):
  54. clsid, score = int(boxes[i]["cls_id"]), boxes[i]["score"]
  55. mask = np_masks[i]
  56. if clsid not in clsid2color:
  57. color_index = i % len(color_list)
  58. clsid2color[clsid] = color_list[color_index]
  59. color_mask = clsid2color[clsid]
  60. for c in range(3):
  61. color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
  62. idx = np.nonzero(mask)
  63. color_mask = np.array(color_mask)
  64. im[idx[0], idx[1], :] *= 1.0 - alpha
  65. im[idx[0], idx[1], :] += alpha * color_mask
  66. return Image.fromarray(im.astype("uint8"))
  67. class InstanceSegResult(BaseResult):
  68. """Save Result Transform"""
  69. def __init__(self, data):
  70. super().__init__(data)
  71. # We use pillow backend to save both numpy arrays and PIL Image objects
  72. self._img_reader.set_backend("pillow")
  73. self._img_writer.set_backend("pillow")
  74. def _get_res_img(self):
  75. """apply"""
  76. boxes = np.array(self["boxes"])
  77. masks = self["masks"]
  78. img_path = self["img_path"]
  79. file_name = os.path.basename(img_path)
  80. image = self._img_reader.read(img_path)
  81. ori_img_size = list(image.size)[::-1]
  82. image = draw_mask(image, boxes, masks, ori_img_size)
  83. image = draw_box(image, boxes)
  84. return image