Spaces:
Build error
Build error
File size: 3,698 Bytes
daf0288 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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
|