diff --git a/utils/plots.py b/utils/plots.py index 3ec7935..94f46a9 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -15,7 +15,7 @@ import pandas as pd import seaborn as sns import torch import yaml -from PIL import Image, ImageDraw +from PIL import Image, ImageDraw, ImageFont from scipy.signal import butter, filtfilt from utils.general import xywh2xyxy, xyxy2xywh @@ -68,6 +68,20 @@ def plot_one_box(x, img, color=None, label=None, line_thickness=None): cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) +def plot_one_box_PIL(box, img, color=None, label=None, line_thickness=None): + img = Image.fromarray(img) + draw = ImageDraw.Draw(img) + line_thickness = line_thickness or max(int(min(img.size) / 200), 2) + draw.rectangle(box, width=line_thickness, outline=tuple(color)) # plot + if label: + fontsize = max(round(max(img.size) / 40), 12) + font = ImageFont.truetype("Arial.ttf", fontsize) + txt_width, txt_height = font.getsize(label) + draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=tuple(color)) + draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font) + return np.asarray(img) + + def plot_wh_methods(): # from utils.plots import *; plot_wh_methods() # Compares the two methods for width-height anchor multiplication # https://github.com/ultralytics/yolov3/issues/168