from collections.abc import Sequence import numpy as np import cv2 from PIL import Image, ImageDraw, ImageFont class Colors: """ Ultralytics default color palette https://ultralytics.com/. This class provides methods to work with the Ultralytics color palette, including converting hex color codes to RGB values. Attributes: palette (list of tuple): List of RGB color values. n (int): The number of colors in the palette. pose_palette (np.ndarray): A specific color palette array with dtype np.uint8. """ def __init__(self): """Initialize colors as 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) self.pose_palette = np.array( [ [255, 128, 0], [255, 153, 51], [255, 178, 102], [230, 230, 0], [255, 153, 255], [153, 204, 255], [255, 102, 255], [255, 51, 255], [102, 178, 255], [51, 153, 255], [255, 153, 153], [255, 102, 102], [255, 51, 51], [153, 255, 153], [102, 255, 102], [51, 255, 51], [0, 255, 0], [0, 0, 255], [255, 0, 0], [255, 255, 255], ], dtype=np.uint8, ) def __call__(self, i, bgr=False): """Converts hex color codes to RGB values.""" c = self.palette[int(i) % self.n] return (c[2], c[1], c[0]) if bgr else c @staticmethod def hex2rgb(h): """Converts hex color codes to RGB values (i.e. default PIL order).""" return tuple(int(h[1 + i: 1 + i + 2], 16) for i in (0, 2, 4)) colors = Colors() class Annotator: def __init__(self, im, line_width=None, font_size=None, font="Arial.ttf", pil=False, example="abc"): self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) self.im = Image.fromarray(im) self.draw = ImageDraw.Draw(self.im) font = 'static/font/Arial.Unicode.ttf' size = font_size or max(round(sum(self.im.size) / 2 * 0.035), 12) self.font = ImageFont.truetype(str(font), size) if not hasattr(self.font, 'getsize'): self.font.getsize = lambda x: self.font.getbbox(x)[2:4] self.tf = max(self.lw - 1, 1) # font thickness self.sf = self.lw / 3 # font scale def box_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255), rotated=False): if not isinstance(box, Sequence): box = box.tolist() if rotated: p1 = box[0] # NOTE: PIL-version polygon needs tuple type. self.draw.polygon([tuple(b) for b in box], width=self.lw, outline=color) else: p1 = (box[0], box[1]) self.draw.rectangle(box, width=self.lw, outline=color) # box if label: w, h = self.font.getsize(label) # text width, height outside = p1[1] - h >= 0 # label fits outside box self.draw.rectangle( (p1[0], p1[1] - h if outside else p1[1], p1[0] + w + 1, p1[1] + 1 if outside else p1[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((p1[0], p1[1] - h if outside else p1[1]), label, fill=txt_color, font=self.font) def result(self): """Return annotated image as array.""" return np.asarray(self.im)