Spaces:
Build error
Build error
| 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 | |