You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
182 lines
8.6 KiB
182 lines
8.6 KiB
from pathlib import Path
|
|
from urllib.error import URLError
|
|
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
from ultralytics.yolo.utils import CONFIG_DIR, FONT
|
|
|
|
from .checks import check_font, check_requirements, is_ascii
|
|
from .files import increment_path
|
|
from .ops import clip_coords, scale_image, xywh2xyxy, xyxy2xywh
|
|
|
|
|
|
class Colors:
|
|
# Ultralytics color palette https://ultralytics.com/
|
|
def __init__(self):
|
|
# hex = matplotlib.colors.TABLEAU_COLORS.values()
|
|
hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
|
|
'2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
|
|
self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
|
|
self.n = len(self.palette)
|
|
|
|
def __call__(self, i, bgr=False):
|
|
c = self.palette[int(i) % self.n]
|
|
return (c[2], c[1], c[0]) if bgr else c
|
|
|
|
@staticmethod
|
|
def hex2rgb(h): # rgb order (PIL)
|
|
return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
|
|
|
|
|
|
colors = Colors() # create instance for 'from utils.plots import colors'
|
|
|
|
|
|
class Annotator:
|
|
# YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
|
|
def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):
|
|
assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.'
|
|
non_ascii = not is_ascii(example) # non-latin labels, i.e. asian, arabic, cyrillic
|
|
self.pil = pil or non_ascii
|
|
if self.pil: # use PIL
|
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
self.draw = ImageDraw.Draw(self.im)
|
|
self.font = check_pil_font(font='Arial.Unicode.ttf' if non_ascii else font,
|
|
size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12))
|
|
else: # use cv2
|
|
self.im = im
|
|
self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width
|
|
|
|
def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)):
|
|
# Add one xyxy box to image with label
|
|
if self.pil or not is_ascii(label):
|
|
self.draw.rectangle(box, width=self.lw, outline=color) # box
|
|
if label:
|
|
w, h = self.font.getsize(label) # text width, height
|
|
outside = box[1] - h >= 0 # label fits outside box
|
|
self.draw.rectangle(
|
|
(box[0], box[1] - h if outside else box[1], box[0] + w + 1,
|
|
box[1] + 1 if outside else box[1] + h + 1),
|
|
fill=color,
|
|
)
|
|
# self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0
|
|
self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font)
|
|
else: # cv2
|
|
p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
|
|
cv2.rectangle(self.im, p1, p2, color, thickness=self.lw, lineType=cv2.LINE_AA)
|
|
if label:
|
|
tf = max(self.lw - 1, 1) # font thickness
|
|
w, h = cv2.getTextSize(label, 0, fontScale=self.lw / 3, thickness=tf)[0] # text width, height
|
|
outside = p1[1] - h >= 3
|
|
p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3
|
|
cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled
|
|
cv2.putText(self.im,
|
|
label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2),
|
|
0,
|
|
self.lw / 3,
|
|
txt_color,
|
|
thickness=tf,
|
|
lineType=cv2.LINE_AA)
|
|
|
|
def masks(self, masks, colors, im_gpu=None, alpha=0.5):
|
|
"""Plot masks at once.
|
|
Args:
|
|
masks (tensor): predicted masks on cuda, shape: [n, h, w]
|
|
colors (List[List[Int]]): colors for predicted masks, [[r, g, b] * n]
|
|
im_gpu (tensor): img is in cuda, shape: [3, h, w], range: [0, 1]
|
|
alpha (float): mask transparency: 0.0 fully transparent, 1.0 opaque
|
|
"""
|
|
if self.pil:
|
|
# convert to numpy first
|
|
self.im = np.asarray(self.im).copy()
|
|
if im_gpu is None:
|
|
# Add multiple masks of shape(h,w,n) with colors list([r,g,b], [r,g,b], ...)
|
|
if len(masks) == 0:
|
|
return
|
|
if isinstance(masks, torch.Tensor):
|
|
masks = torch.as_tensor(masks, dtype=torch.uint8)
|
|
masks = masks.permute(1, 2, 0).contiguous()
|
|
masks = masks.cpu().numpy()
|
|
# masks = np.ascontiguousarray(masks.transpose(1, 2, 0))
|
|
masks = scale_image(masks.shape[:2], masks, self.im.shape)
|
|
masks = np.asarray(masks, dtype=np.float32)
|
|
colors = np.asarray(colors, dtype=np.float32) # shape(n,3)
|
|
s = masks.sum(2, keepdims=True).clip(0, 1) # add all masks together
|
|
masks = (masks @ colors).clip(0, 255) # (h,w,n) @ (n,3) = (h,w,3)
|
|
self.im[:] = masks * alpha + self.im * (1 - s * alpha)
|
|
else:
|
|
if len(masks) == 0:
|
|
self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
|
|
colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
|
|
colors = colors[:, None, None] # shape(n,1,1,3)
|
|
masks = masks.unsqueeze(3) # shape(n,h,w,1)
|
|
masks_color = masks * (colors * alpha) # shape(n,h,w,3)
|
|
|
|
inv_alph_masks = (1 - masks * alpha).cumprod(0) # shape(n,h,w,1)
|
|
mcs = (masks_color * inv_alph_masks).sum(0) * 2 # mask color summand shape(n,h,w,3)
|
|
|
|
im_gpu = im_gpu.flip(dims=[0]) # flip channel
|
|
im_gpu = im_gpu.permute(1, 2, 0).contiguous() # shape(h,w,3)
|
|
im_gpu = im_gpu * inv_alph_masks[-1] + mcs
|
|
im_mask = (im_gpu * 255).byte().cpu().numpy()
|
|
self.im[:] = scale_image(im_gpu.shape, im_mask, self.im.shape)
|
|
if self.pil:
|
|
# convert im back to PIL and update draw
|
|
self.fromarray(self.im)
|
|
|
|
def rectangle(self, xy, fill=None, outline=None, width=1):
|
|
# Add rectangle to image (PIL-only)
|
|
self.draw.rectangle(xy, fill, outline, width)
|
|
|
|
def text(self, xy, text, txt_color=(255, 255, 255), anchor='top'):
|
|
# Add text to image (PIL-only)
|
|
if anchor == 'bottom': # start y from font bottom
|
|
w, h = self.font.getsize(text) # text width, height
|
|
xy[1] += 1 - h
|
|
self.draw.text(xy, text, fill=txt_color, font=self.font)
|
|
|
|
def fromarray(self, im):
|
|
# Update self.im from a numpy array
|
|
self.im = im if isinstance(im, Image.Image) else Image.fromarray(im)
|
|
self.draw = ImageDraw.Draw(self.im)
|
|
|
|
def result(self):
|
|
# Return annotated image as array
|
|
return np.asarray(self.im)
|
|
|
|
|
|
def check_pil_font(font=FONT, size=10):
|
|
# Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary
|
|
font = Path(font)
|
|
font = font if font.exists() else (CONFIG_DIR / font.name)
|
|
try:
|
|
return ImageFont.truetype(str(font) if font.exists() else font.name, size)
|
|
except Exception: # download if missing
|
|
try:
|
|
check_font(font)
|
|
return ImageFont.truetype(str(font), size)
|
|
except TypeError:
|
|
check_requirements('Pillow>=8.4.0') # known issue https://github.com/ultralytics/yolov5/issues/5374
|
|
except URLError: # not online
|
|
return ImageFont.load_default()
|
|
|
|
|
|
def save_one_box(xyxy, im, file=Path('im.jpg'), gain=1.02, pad=10, square=False, BGR=False, save=True):
|
|
# Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop
|
|
xyxy = torch.tensor(xyxy).view(-1, 4)
|
|
b = xyxy2xywh(xyxy) # boxes
|
|
if square:
|
|
b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square
|
|
b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad
|
|
xyxy = xywh2xyxy(b).long()
|
|
clip_coords(xyxy, im.shape)
|
|
crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)]
|
|
if save:
|
|
file.parent.mkdir(parents=True, exist_ok=True) # make directory
|
|
f = str(increment_path(file).with_suffix('.jpg'))
|
|
# cv2.imwrite(f, crop) # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue
|
|
Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0) # save RGB
|
|
return crop
|