import os from typing import List import logging import PIL from PIL import Image from PIL import ImageDraw from .annotation import Annotation def draw_box(im:Image.Image, result, lables, threshold=0.5): im = im.copy() draw_thickness = min(im.size) // 320 draw = ImageDraw.Draw(im) color_list = get_color_map_list(len(lables)) clsid2color = {n.lower():color_list[i] for i,n in enumerate(lables)} result = [r for r in result if r["score"] >= threshold] for dt in result: color = tuple(clsid2color[dt["type"]]) xmin, ymin, xmax, ymax = dt["bbox"] #Draws a line forming a rectangle around the detected object. draw.line( [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)], width=draw_thickness, fill=color) #Prepares the text for the label (class type and score). # draw label text = "{} {:.4f}".format(dt["type"], dt["score"]) #Computes the size of the text using imagedraw_textsize_c. tw, th = imagedraw_textsize_c(draw, text) #Draws a filled rectangle for the text background. draw.rectangle( [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color) #Draws the text on top of the rectangle. draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255)) return im def draw_only_box(im:Image.Image, result): im = im.copy() draw_thickness = min(im.size) // 400 draw = ImageDraw.Draw(im) result = [r for r in result] for dt in result: xmin, ymin, xmax, ymax = dt xmin = int(xmin) ymin = int(ymin) xmax = int(xmax) ymax = int(ymax) draw.line( [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)], width=draw_thickness, fill="red") return im def draw_box_with_text(im:Image.Image, result:List[Annotation], threshold=0.5): im = im.copy() draw_thickness = min(im.size) // 320 draw = ImageDraw.Draw(im) result = [r for r in result if r.score >= threshold] for dt in result: xmin, ymin, xmax, ymax = dt.box draw.line( [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), (xmin, ymin)], width=draw_thickness, fill="red") # draw label text = "{:.4f}".format(dt.score) tw, th = imagedraw_textsize_c(draw, text) draw.rectangle( [(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill="green") draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255)) return im def get_color_map_list(num_classes): """ Args: num_classes (int): number of class Returns: color_map (list): RGB color list """ color_map = num_classes * [0, 0, 0] for i in range(0, num_classes): j = 0 lab = i while lab: color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j)) color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j)) color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j)) j += 1 lab >>= 3 color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)] return color_map def imagedraw_textsize_c(draw, text): if int(PIL.__version__.split('.')[0]) < 10: tw, th = draw.textsize(text) else: left, top, right, bottom = draw.textbbox((0, 0), text) tw, th = right - left, bottom - top return tw, th