| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102 |
- # 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_colormap, font_colormap
- from .base import BaseResult
- def draw_box(img, np_boxes, labels):
- """
- Args:
- img (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]
- labels (list): labels:['class1', ..., 'classn']
- Returns:
- img (PIL.Image.Image): visualized image
- """
- font_size = int(0.024 * int(img.width)) + 2
- font = ImageFont.truetype(PINGFANG_FONT_FILE_PATH, font_size, encoding="utf-8")
- draw_thickness = int(max(img.size) * 0.005)
- draw = ImageDraw.Draw(img)
- clsid2color = {}
- catid2fontcolor = {}
- color_list = get_colormap(rgb=True)
- for i, dt in enumerate(np_boxes):
- clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
- if clsid not in clsid2color:
- color_index = i % len(color_list)
- clsid2color[clsid] = color_list[color_index]
- catid2fontcolor[clsid] = font_colormap(color_index)
- color = tuple(clsid2color[clsid])
- font_color = tuple(catid2fontcolor[clsid])
- xmin, ymin, xmax, ymax = bbox
- # draw bbox
- draw.line(
- [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)],
- width=draw_thickness,
- fill=color,
- )
- # draw label
- text = "{} {:.2f}".format(labels[clsid], score)
- if tuple(map(int, PIL.__version__.split("."))) <= (10, 0, 0):
- tw, th = draw.textsize(text, font=font)
- else:
- left, top, right, bottom = draw.textbbox((0, 0), text, font)
- tw, th = right - left, bottom - top
- if ymin < th:
- draw.rectangle([(xmin, ymin), (xmin + tw + 4, ymin + th + 1)], fill=color)
- draw.text((xmin + 2, ymin - 2), text, fill=font_color, font=font)
- else:
- draw.rectangle([(xmin, ymin - th), (xmin + tw + 4, ymin + 1)], fill=color)
- draw.text((xmin + 2, ymin - th - 2), text, fill=font_color, font=font)
- return img
- class DetResult(BaseResult):
- """Save Result Transform"""
- def __init__(self, data):
- super().__init__(data)
- self.data = 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"]
- img_path = self["img_path"]
- labels = self.data["labels"]
- file_name = os.path.basename(img_path)
- image = self._img_reader.read(img_path)
- image = draw_box(image, boxes, labels=labels)
- return image
|