# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import os import numpy as np import math import PIL from PIL import Image, ImageDraw, ImageFont from ...utils import logging from ...utils.fonts import PINGFANG_FONT_FILE_PATH from ..utils.io import ImageWriter, ImageReader from ..utils.color_map import get_color_map_list, font_colormap from .base import BaseResult from .det import draw_box def draw_mask(im, np_boxes, np_masks, labels): """ Args: im (PIL.Image.Image): PIL image np_boxes (np.ndarray): shape:[N,6], N: number of box, matix element:[class, score, x_min, y_min, x_max, y_max] np_masks (np.ndarray): shape:[N, im_h, im_w] labels (list): labels:['class1', ..., 'classn'] Returns: im (PIL.Image.Image): visualized image """ color_list = get_color_map_list(len(labels)) w_ratio = 0.4 alpha = 0.7 im = np.array(im).astype("float32") clsid2color = {} im_h, im_w = im.shape[:2] np_masks = np_masks[:, :im_h, :im_w] for i in range(len(np_masks)): clsid, score = int(np_boxes[i][0]), np_boxes[i][1] mask = np_masks[i] if clsid not in clsid2color: clsid2color[clsid] = color_list[clsid] color_mask = clsid2color[clsid] for c in range(3): color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255 idx = np.nonzero(mask) color_mask = np.array(color_mask) im[idx[0], idx[1], :] *= 1.0 - alpha im[idx[0], idx[1], :] += alpha * color_mask return Image.fromarray(im.astype("uint8")) class InstanceSegResult(BaseResult): """Save Result Transform""" def __init__(self, data): super().__init__(data) # We use pillow backend to save both numpy arrays and PIL Image objects self._img_reader.set_backend("pillow") self._img_writer.set_backend("pillow") def _get_res_img(self): """apply""" boxes = self["boxes"] masks = self["masks"] img_path = self["img_path"] labels = self["labels"] file_name = os.path.basename(img_path) image = self._img_reader.read(img_path) image = draw_mask(image, boxes, masks, labels) image = draw_box(image, boxes, labels=labels) return image