Source code for draugr.opencv_utilities.drawing.draw_boxes

#!/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()