509 lines
23 KiB
Python
509 lines
23 KiB
Python
|
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
# ==============================================================================
|
||
|
|
|
||
|
|
"""A set of functions that are used for visualization.
|
||
|
|
These functions often receive an image, perform some visualization on the image.
|
||
|
|
The functions do not return a value, instead they modify the image itself.
|
||
|
|
"""
|
||
|
|
import abc
|
||
|
|
import collections
|
||
|
|
# Set headless-friendly backend.
|
||
|
|
import matplotlib;
|
||
|
|
|
||
|
|
matplotlib.use('Agg') # pylint: disable=multiple-statements
|
||
|
|
import matplotlib.pyplot as plt # pylint: disable=g-import-not-at-top
|
||
|
|
import numpy as np
|
||
|
|
import PIL.Image as Image
|
||
|
|
import PIL.ImageColor as ImageColor
|
||
|
|
import PIL.ImageDraw as ImageDraw
|
||
|
|
import PIL.ImageFont as ImageFont
|
||
|
|
import six
|
||
|
|
from six.moves import range
|
||
|
|
from six.moves import zip
|
||
|
|
import tensorflow as tf
|
||
|
|
|
||
|
|
_TITLE_LEFT_MARGIN = 10
|
||
|
|
_TITLE_TOP_MARGIN = 10
|
||
|
|
STANDARD_COLORS = [
|
||
|
|
'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure', 'Beige', 'Bisque',
|
||
|
|
'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite',
|
||
|
|
'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson', 'Cyan',
|
||
|
|
'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange',
|
||
|
|
'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet',
|
||
|
|
'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite',
|
||
|
|
'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod',
|
||
|
|
'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki',
|
||
|
|
'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue',
|
||
|
|
'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey',
|
||
|
|
'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue',
|
||
|
|
'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime',
|
||
|
|
'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid',
|
||
|
|
'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen',
|
||
|
|
'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin',
|
||
|
|
'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed',
|
||
|
|
'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed',
|
||
|
|
'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple',
|
||
|
|
'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'Green', 'SandyBrown',
|
||
|
|
'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue',
|
||
|
|
'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow',
|
||
|
|
'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Violet', 'Wheat', 'White',
|
||
|
|
'WhiteSmoke', 'Yellow', 'YellowGreen'
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def _get_multiplier_for_color_randomness():
|
||
|
|
"""Returns a multiplier to get semi-random colors from successive indices.
|
||
|
|
This function computes a prime number, p, in the range [2, 17] that:
|
||
|
|
- is closest to len(STANDARD_COLORS) / 10
|
||
|
|
- does not divide len(STANDARD_COLORS)
|
||
|
|
If no prime numbers in that range satisfy the constraints, p is returned as 1.
|
||
|
|
Once p is established, it can be used as a multiplier to select
|
||
|
|
non-consecutive colors from STANDARD_COLORS:
|
||
|
|
colors = [(p * i) % len(STANDARD_COLORS) for i in range(20)]
|
||
|
|
"""
|
||
|
|
num_colors = len(STANDARD_COLORS)
|
||
|
|
prime_candidates = [5, 7, 11, 13, 17]
|
||
|
|
|
||
|
|
# Remove all prime candidates that divide the number of colors.
|
||
|
|
prime_candidates = [p for p in prime_candidates if num_colors % p]
|
||
|
|
if not prime_candidates:
|
||
|
|
return 1
|
||
|
|
|
||
|
|
# Return the closest prime number to num_colors / 10.
|
||
|
|
abs_distance = [np.abs(num_colors / 10. - p) for p in prime_candidates]
|
||
|
|
num_candidates = len(abs_distance)
|
||
|
|
inds = [i for _, i in sorted(zip(abs_distance, range(num_candidates)))]
|
||
|
|
return prime_candidates[inds[0]]
|
||
|
|
|
||
|
|
|
||
|
|
def draw_bounding_box_on_image_array(image,
|
||
|
|
ymin,
|
||
|
|
xmin,
|
||
|
|
ymax,
|
||
|
|
xmax,
|
||
|
|
color='red',
|
||
|
|
thickness=4,
|
||
|
|
display_str_list=(),
|
||
|
|
use_normalized_coordinates=True):
|
||
|
|
"""Adds a bounding box to an image (numpy array).
|
||
|
|
Bounding box coordinates can be specified in either absolute (pixel) or
|
||
|
|
normalized coordinates by setting the use_normalized_coordinates argument.
|
||
|
|
Args:
|
||
|
|
image: a numpy array with shape [height, width, 3].
|
||
|
|
ymin: ymin of bounding box.
|
||
|
|
xmin: xmin of bounding box.
|
||
|
|
ymax: ymax of bounding box.
|
||
|
|
xmax: xmax of bounding box.
|
||
|
|
color: color to draw bounding box. Default is red.
|
||
|
|
thickness: line thickness. Default value is 4.
|
||
|
|
display_str_list: list of strings to display in box
|
||
|
|
(each to be shown on its own line).
|
||
|
|
use_normalized_coordinates: If True (default), treat coordinates
|
||
|
|
ymin, xmin, ymax, xmax as relative to the image. Otherwise treat
|
||
|
|
coordinates as absolute.
|
||
|
|
"""
|
||
|
|
image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
|
||
|
|
draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,
|
||
|
|
thickness, display_str_list,
|
||
|
|
use_normalized_coordinates)
|
||
|
|
np.copyto(image, np.array(image_pil))
|
||
|
|
|
||
|
|
|
||
|
|
def draw_bounding_box_on_image(image,
|
||
|
|
ymin,
|
||
|
|
xmin,
|
||
|
|
ymax,
|
||
|
|
xmax,
|
||
|
|
color='red',
|
||
|
|
thickness=4,
|
||
|
|
display_str_list=(),
|
||
|
|
use_normalized_coordinates=True):
|
||
|
|
"""Adds a bounding box to an image.
|
||
|
|
Bounding box coordinates can be specified in either absolute (pixel) or
|
||
|
|
normalized coordinates by setting the use_normalized_coordinates argument.
|
||
|
|
Each string in display_str_list is displayed on a separate line above the
|
||
|
|
bounding box in black text on a rectangle filled with the input 'color'.
|
||
|
|
If the top of the bounding box extends to the edge of the image, the strings
|
||
|
|
are displayed below the bounding box.
|
||
|
|
Args:
|
||
|
|
image: a PIL.Image object.
|
||
|
|
ymin: ymin of bounding box.
|
||
|
|
xmin: xmin of bounding box.
|
||
|
|
ymax: ymax of bounding box.
|
||
|
|
xmax: xmax of bounding box.
|
||
|
|
color: color to draw bounding box. Default is red.
|
||
|
|
thickness: line thickness. Default value is 4.
|
||
|
|
display_str_list: list of strings to display in box
|
||
|
|
(each to be shown on its own line).
|
||
|
|
use_normalized_coordinates: If True (default), treat coordinates
|
||
|
|
ymin, xmin, ymax, xmax as relative to the image. Otherwise treat
|
||
|
|
coordinates as absolute.
|
||
|
|
"""
|
||
|
|
draw = ImageDraw.Draw(image)
|
||
|
|
im_width, im_height = image.size
|
||
|
|
if use_normalized_coordinates:
|
||
|
|
(left, right, top, bottom) = (xmin * im_width, xmax * im_width,
|
||
|
|
ymin * im_height, ymax * im_height)
|
||
|
|
else:
|
||
|
|
(left, right, top, bottom) = (xmin, xmax, ymin, ymax)
|
||
|
|
if thickness > 0:
|
||
|
|
draw.line([(left, top), (left, bottom), (right, bottom), (right, top),
|
||
|
|
(left, top)],
|
||
|
|
width=thickness,
|
||
|
|
fill=color)
|
||
|
|
try:
|
||
|
|
font = ImageFont.truetype('arial.ttf', 24)
|
||
|
|
except IOError:
|
||
|
|
font = ImageFont.load_default()
|
||
|
|
|
||
|
|
# If the total height of the display strings added to the top of the bounding
|
||
|
|
# box exceeds the top of the image, stack the strings below the bounding box
|
||
|
|
# instead of above.
|
||
|
|
display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]
|
||
|
|
# Each display_str has a top and bottom margin of 0.05x.
|
||
|
|
total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)
|
||
|
|
|
||
|
|
if top > total_display_str_height:
|
||
|
|
text_bottom = top
|
||
|
|
else:
|
||
|
|
text_bottom = bottom + total_display_str_height
|
||
|
|
# Reverse list and print from bottom to top.
|
||
|
|
for display_str in display_str_list[::-1]:
|
||
|
|
text_width, text_height = font.getsize(display_str)
|
||
|
|
margin = np.ceil(0.05 * text_height)
|
||
|
|
draw.rectangle(
|
||
|
|
[(left, text_bottom - text_height - 2 * margin), (left + text_width,
|
||
|
|
text_bottom)],
|
||
|
|
fill=color)
|
||
|
|
draw.text(
|
||
|
|
(left + margin, text_bottom - text_height - margin),
|
||
|
|
display_str,
|
||
|
|
fill='black',
|
||
|
|
font=font)
|
||
|
|
text_bottom -= text_height - 2 * margin
|
||
|
|
|
||
|
|
|
||
|
|
def draw_keypoints_on_image_array(image,
|
||
|
|
key_points,
|
||
|
|
keypoint_scores=None,
|
||
|
|
min_score_thresh=0.5,
|
||
|
|
color='red',
|
||
|
|
radius=2,
|
||
|
|
use_normalized_coordinates=True,
|
||
|
|
keypoint_edges=None,
|
||
|
|
keypoint_edge_color='green',
|
||
|
|
keypoint_edge_width=2):
|
||
|
|
"""Draws keypoints on an image (numpy array).
|
||
|
|
Args:
|
||
|
|
image: a numpy array with shape [height, width, 3].
|
||
|
|
key_points: a numpy array with shape [num_keypoints, 2].
|
||
|
|
keypoint_scores: a numpy array with shape [num_keypoints]. If provided, only
|
||
|
|
those keypoints with a score above score_threshold will be visualized.
|
||
|
|
min_score_thresh: A scalar indicating the minimum keypoint score required
|
||
|
|
for a keypoint to be visualized. Note that keypoint_scores must be
|
||
|
|
provided for this threshold to take effect.
|
||
|
|
color: color to draw the keypoints with. Default is red.
|
||
|
|
radius: keypoint radius. Default value is 2.
|
||
|
|
use_normalized_coordinates: if True (default), treat keypoint values as
|
||
|
|
relative to the image. Otherwise treat them as absolute.
|
||
|
|
keypoint_edges: A list of tuples with keypoint indices that specify which
|
||
|
|
keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
|
||
|
|
edges from keypoint 0 to 1 and from keypoint 2 to 4.
|
||
|
|
keypoint_edge_color: color to draw the keypoint edges with. Default is red.
|
||
|
|
keypoint_edge_width: width of the edges drawn between keypoints. Default
|
||
|
|
value is 2.
|
||
|
|
"""
|
||
|
|
image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
|
||
|
|
draw_keypoints_on_image(image_pil,
|
||
|
|
key_points,
|
||
|
|
keypoint_scores=keypoint_scores,
|
||
|
|
min_score_thresh=min_score_thresh,
|
||
|
|
color=color,
|
||
|
|
radius=radius,
|
||
|
|
use_normalized_coordinates=use_normalized_coordinates,
|
||
|
|
keypoint_edges=keypoint_edges,
|
||
|
|
keypoint_edge_color=keypoint_edge_color,
|
||
|
|
keypoint_edge_width=keypoint_edge_width)
|
||
|
|
np.copyto(image, np.array(image_pil))
|
||
|
|
|
||
|
|
|
||
|
|
def draw_keypoints_on_image(image,
|
||
|
|
key_points,
|
||
|
|
keypoint_scores=None,
|
||
|
|
min_score_thresh=0.5,
|
||
|
|
color='red',
|
||
|
|
radius=2,
|
||
|
|
use_normalized_coordinates=True,
|
||
|
|
keypoint_edges=None,
|
||
|
|
keypoint_edge_color='green',
|
||
|
|
keypoint_edge_width=2):
|
||
|
|
"""Draws keypoints on an image.
|
||
|
|
Args:
|
||
|
|
image: a PIL.Image object.
|
||
|
|
key_points: a numpy array with shape [num_keypoints, 2].
|
||
|
|
keypoint_scores: a numpy array with shape [num_keypoints].
|
||
|
|
min_score_thresh: a score threshold for visualizing keypoints. Only used if
|
||
|
|
keypoint_scores is provided.
|
||
|
|
color: color to draw the keypoints with. Default is red.
|
||
|
|
radius: keypoint radius. Default value is 2.
|
||
|
|
use_normalized_coordinates: if True (default), treat keypoint values as
|
||
|
|
relative to the image. Otherwise treat them as absolute.
|
||
|
|
keypoint_edges: A list of tuples with keypoint indices that specify which
|
||
|
|
keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
|
||
|
|
edges from keypoint 0 to 1 and from keypoint 2 to 4.
|
||
|
|
keypoint_edge_color: color to draw the keypoint edges with. Default is red.
|
||
|
|
keypoint_edge_width: width of the edges drawn between keypoints. Default
|
||
|
|
value is 2.
|
||
|
|
"""
|
||
|
|
draw = ImageDraw.Draw(image)
|
||
|
|
im_width, im_height = image.size
|
||
|
|
key_points = np.array(key_points)
|
||
|
|
key_points_x = [k[1] for k in key_points]
|
||
|
|
key_points_y = [k[0] for k in key_points]
|
||
|
|
if use_normalized_coordinates:
|
||
|
|
key_points_x = tuple([im_width * x for x in key_points_x])
|
||
|
|
key_points_y = tuple([im_height * y for y in key_points_y])
|
||
|
|
if keypoint_scores is not None:
|
||
|
|
keypoint_scores = np.array(keypoint_scores)
|
||
|
|
valid_kpt = np.greater(keypoint_scores, min_score_thresh)
|
||
|
|
else:
|
||
|
|
valid_kpt = np.where(np.any(np.isnan(key_points), axis=1),
|
||
|
|
np.zeros_like(key_points[:, 0]),
|
||
|
|
np.ones_like(key_points[:, 0]))
|
||
|
|
valid_kpt = [v for v in valid_kpt]
|
||
|
|
|
||
|
|
for keypoint_x, keypoint_y, valid in zip(key_points_x, key_points_y, valid_kpt):
|
||
|
|
if valid:
|
||
|
|
draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
|
||
|
|
(keypoint_x + radius, keypoint_y + radius)],
|
||
|
|
outline=color, fill=color)
|
||
|
|
if keypoint_edges is not None:
|
||
|
|
for keypoint_start, keypoint_end in keypoint_edges:
|
||
|
|
if (keypoint_start < 0 or keypoint_start >= len(key_points) or
|
||
|
|
keypoint_end < 0 or keypoint_end >= len(key_points)):
|
||
|
|
continue
|
||
|
|
if not (valid_kpt[keypoint_start] and valid_kpt[keypoint_end]):
|
||
|
|
continue
|
||
|
|
edge_coordinates = [
|
||
|
|
key_points_x[keypoint_start], key_points_y[keypoint_start],
|
||
|
|
key_points_x[keypoint_end], key_points_y[keypoint_end]
|
||
|
|
]
|
||
|
|
draw.line(
|
||
|
|
edge_coordinates, fill=keypoint_edge_color, width=keypoint_edge_width)
|
||
|
|
|
||
|
|
|
||
|
|
def draw_mask_on_image_array(image, mask, color='red', alpha=0.4):
|
||
|
|
"""Draws mask on an image.
|
||
|
|
Args:
|
||
|
|
image: uint8 numpy array with shape (img_height, img_height, 3)
|
||
|
|
mask: a uint8 numpy array of shape (img_height, img_height) with
|
||
|
|
values between either 0 or 1.
|
||
|
|
color: color to draw the keypoints with. Default is red.
|
||
|
|
alpha: transparency value between 0 and 1. (default: 0.4)
|
||
|
|
Raises:
|
||
|
|
ValueError: On incorrect data type for image or masks.
|
||
|
|
"""
|
||
|
|
if image.dtype != np.uint8:
|
||
|
|
raise ValueError('`image` not of type np.uint8')
|
||
|
|
if mask.dtype != np.uint8:
|
||
|
|
raise ValueError('`mask` not of type np.uint8')
|
||
|
|
if np.any(np.logical_and(mask != 1, mask != 0)):
|
||
|
|
raise ValueError('`mask` elements should be in [0, 1]')
|
||
|
|
if image.shape[:2] != mask.shape:
|
||
|
|
raise ValueError('The image has spatial dimensions %s but the mask has '
|
||
|
|
'dimensions %s' % (image.shape[:2], mask.shape))
|
||
|
|
rgb = ImageColor.getrgb(color)
|
||
|
|
pil_image = Image.fromarray(image)
|
||
|
|
|
||
|
|
solid_color = np.expand_dims(
|
||
|
|
np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])
|
||
|
|
pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert('RGBA')
|
||
|
|
pil_mask = Image.fromarray(np.uint8(255.0 * alpha * mask)).convert('L')
|
||
|
|
pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)
|
||
|
|
np.copyto(image, np.array(pil_image.convert('RGB')))
|
||
|
|
|
||
|
|
|
||
|
|
def visualize_boxes_and_labels_on_image_array(
|
||
|
|
image,
|
||
|
|
boxes,
|
||
|
|
classes,
|
||
|
|
scores,
|
||
|
|
category_index,
|
||
|
|
instance_masks=None,
|
||
|
|
instance_boundaries=None,
|
||
|
|
keypoints=None,
|
||
|
|
keypoint_scores=None,
|
||
|
|
keypoint_edges=None,
|
||
|
|
track_ids=None,
|
||
|
|
use_normalized_coordinates=False,
|
||
|
|
max_boxes_to_draw=20,
|
||
|
|
min_score_thresh=.5,
|
||
|
|
agnostic_mode=False,
|
||
|
|
line_thickness=4,
|
||
|
|
groundtruth_box_visualization_color='black',
|
||
|
|
skip_boxes=False,
|
||
|
|
skip_scores=False,
|
||
|
|
skip_labels=False,
|
||
|
|
skip_track_ids=False):
|
||
|
|
"""Overlay labeled boxes on an image with formatted scores and label names.
|
||
|
|
This function groups boxes that correspond to the same location
|
||
|
|
and creates a display string for each detection and overlays these
|
||
|
|
on the image. Note that this function modifies the image in place, and returns
|
||
|
|
that same image.
|
||
|
|
Args:
|
||
|
|
image: uint8 numpy array with shape (img_height, img_width, 3)
|
||
|
|
boxes: a numpy array of shape [N, 4]
|
||
|
|
classes: a numpy array of shape [N]. Note that class indices are 1-based,
|
||
|
|
and match the keys in the label map.
|
||
|
|
scores: a numpy array of shape [N] or None. If scores=None, then
|
||
|
|
this function assumes that the boxes to be plotted are groundtruth
|
||
|
|
boxes and plot all boxes as black with no classes or scores.
|
||
|
|
category_index: a dict containing category dictionaries (each holding
|
||
|
|
category index `id` and category name `name`) keyed by category indices.
|
||
|
|
instance_masks: a numpy array of shape [N, image_height, image_width] with
|
||
|
|
values ranging between 0 and 1, can be None.
|
||
|
|
instance_boundaries: a numpy array of shape [N, image_height, image_width]
|
||
|
|
with values ranging between 0 and 1, can be None.
|
||
|
|
keypoints: a numpy array of shape [N, num_keypoints, 2], can
|
||
|
|
be None.
|
||
|
|
keypoint_scores: a numpy array of shape [N, num_keypoints], can be None.
|
||
|
|
keypoint_edges: A list of tuples with keypoint indices that specify which
|
||
|
|
keypoints should be connected by an edge, e.g. [(0, 1), (2, 4)] draws
|
||
|
|
edges from keypoint 0 to 1 and from keypoint 2 to 4.
|
||
|
|
track_ids: a numpy array of shape [N] with unique track ids. If provided,
|
||
|
|
color-coding of boxes will be determined by these ids, and not the class
|
||
|
|
indices.
|
||
|
|
use_normalized_coordinates: whether boxes is to be interpreted as
|
||
|
|
normalized coordinates or not.
|
||
|
|
max_boxes_to_draw: maximum number of boxes to visualize. If None, draw
|
||
|
|
all boxes.
|
||
|
|
min_score_thresh: minimum score threshold for a box or keypoint to be
|
||
|
|
visualized.
|
||
|
|
agnostic_mode: boolean (default: False) controlling whether to evaluate in
|
||
|
|
class-agnostic mode or not. This mode will display scores but ignore
|
||
|
|
classes.
|
||
|
|
line_thickness: integer (default: 4) controlling line width of the boxes.
|
||
|
|
groundtruth_box_visualization_color: box color for visualizing groundtruth
|
||
|
|
boxes
|
||
|
|
skip_boxes: whether to skip the drawing of bounding boxes.
|
||
|
|
skip_scores: whether to skip score when drawing a single detection
|
||
|
|
skip_labels: whether to skip label when drawing a single detection
|
||
|
|
skip_track_ids: whether to skip track id when drawing a single detection
|
||
|
|
Returns:
|
||
|
|
uint8 numpy array with shape (img_height, img_width, 3) with overlaid boxes.
|
||
|
|
"""
|
||
|
|
# Create a display string (and color) for every box location, group any boxes
|
||
|
|
# that correspond to the same location.
|
||
|
|
box_to_display_str_map = collections.defaultdict(list)
|
||
|
|
box_to_color_map = collections.defaultdict(str)
|
||
|
|
box_to_instance_masks_map = {}
|
||
|
|
box_to_instance_boundaries_map = {}
|
||
|
|
box_to_keypoints_map = collections.defaultdict(list)
|
||
|
|
box_to_keypoint_scores_map = collections.defaultdict(list)
|
||
|
|
box_to_track_ids_map = {}
|
||
|
|
is_crack = False
|
||
|
|
if not max_boxes_to_draw:
|
||
|
|
max_boxes_to_draw = boxes.shape[0]
|
||
|
|
for i in range(boxes.shape[0]):
|
||
|
|
if max_boxes_to_draw == len(box_to_color_map):
|
||
|
|
break
|
||
|
|
if scores is None or scores[i] > min_score_thresh:
|
||
|
|
box = tuple(boxes[i].tolist())
|
||
|
|
if instance_masks is not None:
|
||
|
|
box_to_instance_masks_map[box] = instance_masks[i]
|
||
|
|
if instance_boundaries is not None:
|
||
|
|
box_to_instance_boundaries_map[box] = instance_boundaries[i]
|
||
|
|
if keypoints is not None:
|
||
|
|
box_to_keypoints_map[box].extend(keypoints[i])
|
||
|
|
if keypoint_scores is not None:
|
||
|
|
box_to_keypoint_scores_map[box].extend(keypoint_scores[i])
|
||
|
|
if track_ids is not None:
|
||
|
|
box_to_track_ids_map[box] = track_ids[i]
|
||
|
|
if scores is None:
|
||
|
|
box_to_color_map[box] = groundtruth_box_visualization_color
|
||
|
|
else:
|
||
|
|
display_str = ''
|
||
|
|
if not skip_labels:
|
||
|
|
if not agnostic_mode:
|
||
|
|
if classes[i] in six.viewkeys(category_index):
|
||
|
|
class_name = category_index[classes[i]]['name']
|
||
|
|
else:
|
||
|
|
class_name = 'N/A'
|
||
|
|
display_str = str(class_name)
|
||
|
|
if not skip_scores:
|
||
|
|
if not display_str:
|
||
|
|
display_str = '{}%'.format(round(100 * scores[i]))
|
||
|
|
else:
|
||
|
|
display_str = '{}: {}%'.format(display_str, round(100 * scores[i]))
|
||
|
|
if not skip_track_ids and track_ids is not None:
|
||
|
|
if not display_str:
|
||
|
|
display_str = 'ID {}'.format(track_ids[i])
|
||
|
|
else:
|
||
|
|
display_str = '{}: ID {}'.format(display_str, track_ids[i])
|
||
|
|
box_to_display_str_map[box].append(display_str)
|
||
|
|
if agnostic_mode:
|
||
|
|
box_to_color_map[box] = 'DarkOrange'
|
||
|
|
elif track_ids is not None:
|
||
|
|
prime_multipler = _get_multiplier_for_color_randomness()
|
||
|
|
box_to_color_map[box] = STANDARD_COLORS[
|
||
|
|
(prime_multipler * track_ids[i]) % len(STANDARD_COLORS)]
|
||
|
|
else:
|
||
|
|
box_to_color_map[box] = STANDARD_COLORS[
|
||
|
|
classes[i] % len(STANDARD_COLORS)]
|
||
|
|
is_crack = len(box_to_color_map) > 0
|
||
|
|
# Draw all boxes onto image.
|
||
|
|
for box, color in box_to_color_map.items():
|
||
|
|
ymin, xmin, ymax, xmax = box
|
||
|
|
if instance_masks is not None:
|
||
|
|
draw_mask_on_image_array(
|
||
|
|
image,
|
||
|
|
box_to_instance_masks_map[box],
|
||
|
|
color=color
|
||
|
|
)
|
||
|
|
if instance_boundaries is not None:
|
||
|
|
draw_mask_on_image_array(
|
||
|
|
image,
|
||
|
|
box_to_instance_boundaries_map[box],
|
||
|
|
color='red',
|
||
|
|
alpha=1.0
|
||
|
|
)
|
||
|
|
draw_bounding_box_on_image_array(
|
||
|
|
image,
|
||
|
|
ymin,
|
||
|
|
xmin,
|
||
|
|
ymax,
|
||
|
|
xmax,
|
||
|
|
color=color,
|
||
|
|
thickness=0 if skip_boxes else line_thickness,
|
||
|
|
display_str_list=box_to_display_str_map[box],
|
||
|
|
use_normalized_coordinates=use_normalized_coordinates)
|
||
|
|
if keypoints is not None:
|
||
|
|
keypoint_scores_for_box = None
|
||
|
|
if box_to_keypoint_scores_map:
|
||
|
|
keypoint_scores_for_box = box_to_keypoint_scores_map[box]
|
||
|
|
draw_keypoints_on_image_array(
|
||
|
|
image,
|
||
|
|
box_to_keypoints_map[box],
|
||
|
|
keypoint_scores_for_box,
|
||
|
|
min_score_thresh=min_score_thresh,
|
||
|
|
color=color,
|
||
|
|
radius=line_thickness / 2,
|
||
|
|
use_normalized_coordinates=use_normalized_coordinates,
|
||
|
|
keypoint_edges=keypoint_edges,
|
||
|
|
keypoint_edge_color=color,
|
||
|
|
keypoint_edge_width=line_thickness // 2)
|
||
|
|
|
||
|
|
return image, is_crack
|