Spaces:
Runtime error
Runtime error
File size: 1,841 Bytes
684e6f5 |
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 |
import random
from PIL import ImageDraw
def plot_one_box(x, img, color=None, label=None, line_thickness=None):
"""
Helper Functions for Plotting BBoxes
:param x:
:param img:
:param color:
:param label:
:param line_thickness:
:return:
"""
width, height = img.size
tl = line_thickness or round(0.002 * (width + height) / 2) + 1 # line/font thickness
color = color or (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
img_draw = ImageDraw.Draw(img)
img_draw.rectangle((c1[0], c1[1], c2[0], c2[1]), outline=color, width=tl)
if label:
tf = max(tl - 1, 1) # font thickness
x1, y1, x2, y2 = img_draw.textbbox(c1, label, stroke_width=tf)
img_draw.rectangle((x1, y1, x2, y2), fill=color)
img_draw.text((x1, y1), label, fill=(255, 255, 255))
def add_bboxes(pil_img, result, confidence=0.6):
"""
Plotting Bounding Box on img
:param pil_img:
:param result:
:param confidence:
:return:
"""
for box in result.boxes:
[cl] = box.cls.tolist()
[conf] = box.conf.tolist()
if conf < confidence:
continue
[rect] = box.xyxy.tolist()
text = f'{result.names[cl]}: {conf: 0.2f}'
plot_one_box(x=rect, img=pil_img, label=text)
return pil_img
def add_bboxes2(pil_img, result, confidence=0.6):
"""
Plotting Bounding Box on img
:param pil_img:
:param result:
:param confidence:
:return:
"""
for box in result['boxes']:
cl = box['cls']
conf = box['conf']
if conf < confidence:
continue
rect = box['xyxy']
text = f'{cl}: {conf: 0.2f}'
plot_one_box(x=rect, img=pil_img, label=text)
return pil_img
|