Unverified Commit 2da4e7ac authored by Glenn Jocher's avatar Glenn Jocher Committed by GitHub
Browse files

Merge PIL and OpenCV in `plot_one_box(use_pil=False)` (#4416)

* Merge PIL and OpenCV box plotting functions

* Add ASCII check to plot_one_box

* Cleanup

* Cleanup2
parent 63e09fdc
...@@ -156,7 +156,7 @@ def run(weights='yolov5s.pt', # model.pt path(s) ...@@ -156,7 +156,7 @@ def run(weights='yolov5s.pt', # model.pt path(s)
if save_img or save_crop or view_img: # Add bbox to image if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class c = int(cls) # integer class
label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}') label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')
plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=line_thickness) im0 = plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_width=line_thickness)
if save_crop: if save_crop:
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True) save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
......
...@@ -354,7 +354,7 @@ class Detections: ...@@ -354,7 +354,7 @@ class Detections:
if crop: if crop:
save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i]) save_one_box(box, im, file=save_dir / 'crops' / self.names[int(cls)] / self.files[i])
else: # all others else: # all others
plot_one_box(box, im, label=label, color=colors(cls)) im = plot_one_box(box, im, label=label, color=colors(cls))
else: else:
str += '(no detections)' str += '(no detections)'
......
...@@ -110,9 +110,14 @@ def is_pip(): ...@@ -110,9 +110,14 @@ def is_pip():
return 'site-packages' in Path(__file__).absolute().parts return 'site-packages' in Path(__file__).absolute().parts
def is_ascii(str=''):
# Is string composed of all ASCII (no UTF) characters?
return len(str.encode().decode('ascii', 'ignore')) == len(str)
def emojis(str=''): def emojis(str=''):
# Return platform-dependent emoji-safe version of string # Return platform-dependent emoji-safe version of string
return str.encode().decode(encoding='ascii', errors='ignore') if platform.system() == 'Windows' else str return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str
def file_size(file): def file_size(file):
......
# Plotting utils # Plotting utils
import math
from copy import copy from copy import copy
from pathlib import Path from pathlib import Path
import cv2 import cv2
import math
import matplotlib import matplotlib
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import numpy as np
import pandas as pd import pandas as pd
import seaborn as sn import seaborn as sn
import torch import torch
import yaml
from PIL import Image, ImageDraw, ImageFont from PIL import Image, ImageDraw, ImageFont
from utils.general import xywh2xyxy, xyxy2xywh from utils.general import is_ascii, xyxy2xywh, xywh2xyxy
from utils.metrics import fitness from utils.metrics import fitness
# Settings # Settings
...@@ -65,32 +64,31 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5): ...@@ -65,32 +64,31 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5):
return filtfilt(b, a, data) # forward-backward filter return filtfilt(b, a, data) # forward-backward filter
def plot_one_box(x, im, color=(128, 128, 128), label=None, line_thickness=3): def plot_one_box(box, im, color=(128, 128, 128), txt_color=(255, 255, 255), label=None, line_width=3, use_pil=False):
# Plots one bounding box on image 'im' using OpenCV # Plots one xyxy box on image im with label
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.' assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1 # line/font thickness lw = line_width or max(int(min(im.size) / 200), 2) # line width
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(im, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) if use_pil or not is_ascii(label): # use PIL
if label: im = Image.fromarray(im)
tf = max(tl - 1, 1) # font thickness draw = ImageDraw.Draw(im)
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] draw.rectangle(box, width=lw + 1, outline=color) # plot
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 if label:
cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled font = ImageFont.truetype("Arial.ttf", size=max(round(max(im.size) / 40), 12))
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) txt_width, txt_height = font.getsize(label)
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=color)
draw.text((box[0], box[1] - txt_height + 1), label, fill=txt_color, font=font)
def plot_one_box_PIL(box, im, color=(128, 128, 128), label=None, line_thickness=3): return np.asarray(im)
# Plots one bounding box on image 'im' using PIL else: # use OpenCV
im = Image.fromarray(im) c1, c2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
draw = ImageDraw.Draw(im) cv2.rectangle(im, c1, c2, color, thickness=lw, lineType=cv2.LINE_AA)
line_thickness = line_thickness or max(int(min(im.size) / 200), 2) if label:
draw.rectangle(box, width=line_thickness, outline=color) # plot tf = max(lw - 1, 1) # font thickness
if label: txt_width, txt_height = cv2.getTextSize(label, 0, fontScale=lw / 3, thickness=tf)[0]
font = ImageFont.truetype("Arial.ttf", size=max(round(max(im.size) / 40), 12)) c2 = c1[0] + txt_width, c1[1] - txt_height - 3
txt_width, txt_height = font.getsize(label) cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA) # filled
draw.rectangle([box[0], box[1] - txt_height + 4, box[0] + txt_width, box[1]], fill=color) cv2.putText(im, label, (c1[0], c1[1] - 2), 0, lw / 3, txt_color, thickness=tf, lineType=cv2.LINE_AA)
draw.text((box[0], box[1] - txt_height + 1), label, fill=(255, 255, 255), font=font) return im
return np.asarray(im)
def plot_wh_methods(): # from utils.plots import *; plot_wh_methods() def plot_wh_methods(): # from utils.plots import *; plot_wh_methods()
...@@ -180,7 +178,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max ...@@ -180,7 +178,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max
cls = names[cls] if names else cls cls = names[cls] if names else cls
if labels or conf[j] > 0.25: # 0.25 conf thresh if labels or conf[j] > 0.25: # 0.25 conf thresh
label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j]) label = '%s' % cls if labels else '%s %.1f' % (cls, conf[j])
plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl) mosaic = plot_one_box(box, mosaic, label=label, color=color, line_width=tl)
# Draw image filename labels # Draw image filename labels
if paths: if paths:
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment