#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "Christian Heider Nielsen"
__doc__ = r"""
Created on 04/03/2020
"""
from pathlib import Path
from typing import Optional, Sequence, Tuple, Union
import cv2
import numpy
from PIL import Image, ImageDraw, ImageFont
from draugr.opencv_utilities.drawing.opencv_draw import draw_masks
from warg import RGB, compute_color_for_labels
__all__ = ["draw_single_box", "draw_bounding_boxes", "draw_single_rect"]
[docs]def draw_single_box(
image: Image.Image,
x_min: float,
y_min: float,
x_max: float,
y_max: float,
*,
score_color: Tuple = RGB(0, 0, 0),
display_str: str = None,
font_type: ImageFont = ImageFont.load_default(),
outline_color: Tuple = RGB(0, 255, 0),
outline_width: int = 2,
outline_alpha: float = 0.5,
color_fill_score: bool = False,
) -> Image.Image:
"""description"""
if not isinstance(image, ImageDraw.ImageDraw):
assert isinstance(image, Image.Image)
if image.mode != "RGBA":
image = image.convert("RGBA")
draw = ImageDraw.Draw(image, mode="RGBA")
else:
draw = image
left, right, top, bottom = x_min, x_max, y_min, y_max
alpha_color = outline_color + (int(255 * outline_alpha),)
draw.rectangle(
((left, top), (right, bottom)),
outline=outline_color,
fill=alpha_color if color_fill_score else None,
width=outline_width,
)
if display_str:
display_str = str(display_str)
if font_type is None or not isinstance(
font_type,
(ImageFont.FreeTypeFont, ImageFont.ImageFont, ImageFont.TransposedFont),
):
font_type = ImageFont.load_default()
print("Loading default (Better Than Nothing) font")
text_bottom = bottom
# Reverse list and print from bottom to top.
text_width, text_height = font_type.getsize(display_str)
margin = numpy.ceil(0.05 * text_height)
draw.rectangle(
xy=(
(
left + outline_width,
text_bottom - text_height - 2 * margin - outline_width,
),
(left + text_width + outline_width, text_bottom - outline_width),
),
fill=alpha_color,
)
draw.text(
(
left + margin + outline_width,
text_bottom - text_height - margin - outline_width,
),
display_str,
fill=score_color,
font=font_type,
)
return image
[docs]def draw_single_rect(img, rect):
"""
:param img:
:type img:
:param rect:
:type rect:
:return:
:rtype:
"""
# draw_single_box(img, rect)
box = cv2.boxPoints(rect)
box = numpy.int0(box)
cv2.drawContours(img, [box], 0, (255, 255, 255), -1)
return img
[docs]def draw_bounding_boxes(
image: Union[numpy.ndarray, Image.Image],
boxes: numpy.ndarray,
*,
labels: Optional[numpy.ndarray] = None,
scores: Optional[numpy.ndarray] = None,
categories: Optional[Sequence[str]] = None,
outline_width: int = 2,
outline_alpha: float = 0.5,
score_color_fill: bool = False,
score_font: ImageFont = ImageFont.load_default(),
score_format: str = ": {:.2f}",
mode="RGBA", # TODO: make assertion on the channel size
) -> numpy.ndarray:
"""Draw bounding boxes(labels, scores) on image
Args:
image: numpy array image, shape should be (height, width, channel)
boxes: bboxes, shape should be (N, 4), and each row is (xmin, ymin, xmax, ymax), NOT NORMALISED!
labels: labels, shape: (N, )
scores: label scores, shape: (N, )
class_name_map: list or dict, map class id to class name for visualization.
width: box width
alpha: text background alpha
fill: fill box or not
font: text font
score_format: score format
Returns:
An image with information drawn on it."""
boxes = numpy.array(boxes)
num_boxes = boxes.shape[0]
if isinstance(image, Image.Image):
draw_image = image
elif isinstance(image, numpy.ndarray):
draw_image = Image.fromarray(image, mode)
else:
raise AttributeError(f"Unsupported images type {type(image)}")
for i in range(num_boxes):
display_str = ""
color = (0, 255, 0)
if labels is not None:
ith_cat = int(labels[i])
color = compute_color_for_labels(ith_cat)
class_name = categories[ith_cat] if categories is not None else str(ith_cat)
display_str += class_name
if scores is not None:
display_str += f"{score_format.format(scores[i])}"
draw_image = draw_single_box(
image=draw_image,
x_min=boxes[i, 0],
y_min=boxes[i, 1],
x_max=boxes[i, 2],
y_max=boxes[i, 3],
outline_color=color,
outline_width=outline_width,
outline_alpha=outline_alpha,
display_str=display_str,
font_type=score_font,
color_fill_score=score_color_fill,
)
return numpy.array(draw_image, dtype=numpy.uint8)
if __name__ == "__main__":
def a() -> None:
"""
:rtype: None
"""
from matplotlib import pyplot
import pickle
from neodroidvision.data.detection.coco import COCODataset
name = "000000308476"
data_root = Path.home() / "Data" / "Coco"
with open(str(data_root / f"{name}.pickle"), "rb") as f:
data = pickle.load(f)
img = Image.open(str(data_root / f"{name}.jpg"))
img = draw_masks(img, data.masks, labels=data.labels)
img = draw_bounding_boxes(
img,
boxes=data.boxes,
labels=data.labels,
scores=data.scores,
categories=COCODataset.categories,
score_format=": {:.2f}",
)
pyplot.imshow(img)
pyplot.show()
def b() -> None:
"""
:rtype: None
"""
from matplotlib import pyplot
data_root = Path.home() / "Data" / "Datasets" / "PennFudanPed" / "PNGImages"
img = Image.open(str(data_root / f"FudanPed00001.png"))
img_width, _ = img.size
img = draw_bounding_boxes(
img,
boxes=(
(0 * img_width, 0 * img_width, 0.2 * img_width, 0.2 * img_width),
(0.3 * img_width, 0.3 * img_width, 0.5 * img_width, 0.5 * img_width),
(0.1 * img_width, 0.1 * img_width, 0.7 * img_width, 0.7 * img_width),
),
labels=(0, 55, 78),
scores=(0.2, 1.0, 0.4),
categories=(*[f"cat{i}" for i in range(99)],),
score_format=": {:.2f}",
outline_width=5,
score_color_fill=False,
)
pyplot.imshow(img)
pyplot.show()
def c() -> None:
"""
:rtype: None
"""
from matplotlib import pyplot
res = 999
num_channels = 3
img = Image.fromarray(numpy.zeros((res, res, num_channels), dtype=numpy.uint8))
img_width, _ = img.size
img = draw_bounding_boxes(
img,
boxes=(
(0 * img_width, 0 * img_width, 0.2 * img_width, 0.2 * img_width),
(0.3 * img_width, 0.3 * img_width, 0.5 * img_width, 0.5 * img_width),
(0.1 * img_width, 0.1 * img_width, 0.7 * img_width, 0.7 * img_width),
),
labels=(0, 55, 78),
scores=(0.2, 1.0, 0.4),
categories=(*[f"cat{i}" for i in range(99)],),
score_format=": {:.2f}",
outline_width=5,
score_color_fill=False,
)
pyplot.imshow(img)
pyplot.show()
# a()
# b()
c()