#!/usr/bin/env python3
# -*- coding: utf-8 -*-
__author__ = "Christian Heider Nielsen"
__doc__ = r"""
"""
import functools
from pathlib import Path
import numpy
import tensorflow
from PIL.Image import Image, fromarray
from neodroidvision.utilities.visualisation.bounding_box_visualisation import (
_visualize_boxes,
_visualize_boxes_and_keypoints,
_visualize_boxes_and_masks,
_visualize_boxes_and_masks_and_keypoints,
cdf_plot,
hist_plot,
)
[docs]def save_image_array_as_png(image: Image, output_path: Path) -> None:
"""Saves an image (represented as a numpy array) to PNG.
Args:
image: a numpy array with shape [height, width, 3].
output_path: path to which image should be written.
"""
image_pil = fromarray(numpy.uint8(image)).convert("RGB")
with tensorflow.gfile.Open(output_path, "w") as fid:
image_pil.save(fid, "PNG")
[docs]def draw_bounding_boxes_on_image_tensors(
images,
boxes,
classes,
scores,
category_index,
instance_masks=None,
keypoints=None,
max_boxes_to_draw=20,
min_score_thresh=0.2,
line_thickness=2,
):
"""Draws bounding boxes, masks, and keypoints on batch of image tensors.
Args:
:param images: A 4D uint8 image tensor of shape [N, H, W, C].
:param boxes: [N, max_detections, 4] float32 tensor of detection boxes.
:param classes: [N, max_detections] int tensor of detection classes. Note that
classes are 1-indexed.
:param scores: [N, max_detections] float32 tensor of detection scores.
:param category_index: a dict that maps integer ids to category dicts. e.g.
{1: {1: 'dog'}, 2: {2: 'cat'}, ...}
:param instance_masks: A 4D uint8 tensor of shape [N, max_detection, H, W] with
instance masks.
:param keypoints: A 4D float32 tensor of shape [N, max_detection, num_keypoints, 2]
with keypoints.
:param max_boxes_to_draw: Maximum number of boxes to draw on an image. Default 20.
:param min_score_thresh: Minimum score threshold for visualization. Default 0.2.
:param line_thickness:
Returns:
4D image tensor of type uint8, with boxes drawn on top.
"""
visualization_keyword_args = {
"use_normalized_coordinates": True,
"max_boxes_to_draw": max_boxes_to_draw,
"min_score_thresh": min_score_thresh,
"agnostic_mode": False,
"line_thickness": line_thickness,
}
if instance_masks is not None and keypoints is None:
visualize_boxes_fn = functools.partial(
_visualize_boxes_and_masks,
category_index=category_index,
**visualization_keyword_args,
)
elems = [images, boxes, classes, scores, instance_masks]
elif instance_masks is None and keypoints is not None:
visualize_boxes_fn = functools.partial(
_visualize_boxes_and_keypoints,
category_index=category_index,
**visualization_keyword_args,
)
elems = [images, boxes, classes, scores, keypoints]
elif instance_masks is not None and keypoints is not None:
visualize_boxes_fn = functools.partial(
_visualize_boxes_and_masks_and_keypoints,
category_index=category_index,
**visualization_keyword_args,
)
elems = [images, boxes, classes, scores, instance_masks, keypoints]
else:
visualize_boxes_fn = functools.partial(
_visualize_boxes,
category_index=category_index,
**visualization_keyword_args,
)
elems = [images, boxes, classes, scores]
def draw_boxes(image_and_detections):
"""Draws boxes on image."""
image_with_boxes = tensorflow.py_func(
visualize_boxes_fn, image_and_detections, tensorflow.uint8
)
return image_with_boxes
images = tensorflow.map_fn(
draw_boxes, elems, dtype=tensorflow.uint8, back_prop=False
)
return images
[docs]def add_cdf_image_summary(values, name):
"""Adds a tf.summary.image for a CDF plot of the values.
Normalizes `values` such that they sum to 1, plots the cumulative distribution
function and creates a tf image summary.
Args:
values: a 1-D float32 tensor containing the values.
name: name for the image summary."""
tensorflow.summary.image(
name, tensorflow.py_func(cdf_plot, [values], tensorflow.uint8)
)
[docs]def add_hist_image_summary(values, bins, name):
"""Adds a tf.summary.image for a histogram plot of the values.
Plots the histogram of values and creates a tf image summary.
Args:
values: a 1-D float32 tensor containing the values.
bins: bin edges which will be directly passed to numpy.histogram.
name: name for the image summary."""
tensorflow.summary.image(
name, tensorflow.py_func(hist_plot, [values, bins], tensorflow.uint8)
)
