Spaces:
Running
Running
File size: 5,598 Bytes
67a9b5d |
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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
import numpy as np
import PIL
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
from PIL import ImageColor
def _is_legal_color(color):
if color is None:
return True
if isinstance(color, str):
return True
return isinstance(color, (tuple, list)) and len(color) == 3
def _normalize_color(color, pil_mode, swap_rgb=False):
if color is None:
return color
if isinstance(color, str):
color = ImageColor.getrgb(color)
gray = color[0]
if swap_rgb:
color = (color[2], color[1], color[0])
if pil_mode == 'L':
color = gray
return color
def draw_text(image, text, position, color=(255,0,0), font=None, font_size=15):
"""Draws text on given image.
Args:
image (ndarray).
text (str): text to be drawn.
position (Tuple[int, int]): position where to be drawn.
color (List[Union[str, Tuple[int, int, int]]]): text color.
font (str): A filename or file-like object containing a TrueType font. If the file is not found in this
filename, the loader may also search in other directories, such as the `fonts/` directory on Windows
or `/Library/Fonts/`, `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
font_size (int): The requested font size in points.
References:
torchvision.utils.draw_bounding_boxes
"""
if isinstance(image, np.ndarray):
# PIL.Image.fromarray fails with uint16 arrays
# https://github.com/python-pillow/Pillow/issues/1514
if (image.dtype == np.uint16) and (image.ndim != 2):
image = (image / 256).astype(np.uint8)
pil_image = Image.fromarray(image)
elif isinstance(image, PIL.Image.Image):
pil_image = image
else:
raise TypeError('Unsupported image type!')
assert pil_image.mode in ['L', 'RGB', 'RGBA']
assert _is_legal_color(color)
color = _normalize_color(color, pil_image.mode, isinstance(image, np.ndarray))
if font is None:
font_object = ImageFont.load_default()
else:
font_object = ImageFont.truetype(font, size=font_size)
draw = ImageDraw.Draw(pil_image)
draw.text((position[0], position[1]), text,
fill=color, font=font_object)
if isinstance(image, np.ndarray):
return np.asarray(pil_image)
return pil_image
def draw_bounding_boxes(image, boxes, labels=None, colors=None,
fill=False, width=1, font=None, font_size=15):
"""Draws bounding boxes on given image.
Args:
image (ndarray).
boxes (ndarray): ndarray of size (N, 4) containing bounding boxes in (xmin, ymin, xmax, ymax) format.
labels (List[str]): List containing the labels of bounding boxes.
colors (List[Union[str, Tuple[int, int, int]]]): List containing the colors of bounding boxes or labels.
fill (bool): If `True` fills the bounding box with specified color.
width (int): Width of bounding box.
font (str): A filename or file-like object containing a TrueType font. If the file is not found in this
filename, the loader may also search in other directories, such as the `fonts/` directory on Windows
or `/Library/Fonts/`, `/System/Library/Fonts/` and `~/Library/Fonts/` on macOS.
font_size (int): The requested font size in points.
References:
torchvision.utils.draw_bounding_boxes
"""
if isinstance(image, np.ndarray):
# PIL.Image.fromarray fails with uint16 arrays
# https://github.com/python-pillow/Pillow/issues/1514
if (image.dtype == np.uint16) and (image.ndim != 2):
image = (image / 256).astype(np.uint8)
pil_image = Image.fromarray(image)
elif isinstance(image, PIL.Image.Image):
pil_image = image
else:
raise TypeError('Unsupported image type!')
pil_image = pil_image.convert('RGB')
if font is None:
font_object = ImageFont.load_default()
else:
font_object = ImageFont.truetype(font, size=font_size)
if fill:
draw = ImageDraw.Draw(pil_image, "RGBA")
else:
draw = ImageDraw.Draw(pil_image)
for i, bbox in enumerate(boxes):
if colors is None:
color = None
else:
color = colors[i]
assert _is_legal_color(color)
color = _normalize_color(color, pil_image.mode, isinstance(image, np.ndarray))
if fill:
if color is None:
fill_color = (255, 255, 255, 100)
elif isinstance(color, str):
# This will automatically raise Error if rgb cannot be parsed.
fill_color = ImageColor.getrgb(color) + (100,)
elif isinstance(color, tuple):
fill_color = color + (100,)
# the first argument of ImageDraw.rectangle:
# in old version only supports [(x0, y0), (x1, y1)]
# in new version supports either [(x0, y0), (x1, y1)] or [x0, y0, x1, y1]
draw.rectangle([(bbox[0], bbox[1]), (bbox[2], bbox[3])], width=width, outline=color, fill=fill_color)
else:
draw.rectangle([(bbox[0], bbox[1]), (bbox[2], bbox[3])], width=width, outline=color)
if labels is not None:
margin = width + 1
draw.text((bbox[0] + margin, bbox[1] + margin), labels[i], fill=color, font=font_object)
if isinstance(image, np.ndarray):
return np.asarray(pil_image)
return pil_image
|