def testBox2DCornerAxisAligned(self):
# Center at 1., 1., width of 1. and length of 2.
#
# The heading of 0. indicates the heading of the 'long'
# side of the box, so it is longer in the positive y direction.
box = transform_util.Box2D(1.0, 1.0, 2., 1., 0.)
self.assertAllClose(box.corners,
[[0.5, 2.], [1.5, 2.], [1.5, 0.], [0.5, 0.]])
def DrawBoundingBoxOnImage(image,
box,
display_str,
color='red',
thickness=4,
text_loc='BOTTOM'):
"""Draw bounding box on the input image."""
original_image = image
image = Image.fromarray(np.uint8(original_image)).convert('RGB')
draw = ImageDraw.Draw(image)
center_x, center_y, width, height, heading = box
box2d = transform_util.Box2D(center_x, center_y, width, height, heading)
corners = list(box2d.corners.reshape(-1, 2))
points = [tuple(c) for c in corners]
points += [points[0]]
draw.line(points, fill=color, width=thickness)
# Draw heading.
max_dim = max(width, height) / 2.
end_heading_point = (center_x + max_dim * math.cos(heading),
center_y + max_dim * math.sin(heading))
start_heading_point = ((end_heading_point[0] - center_x) / 2 + center_x,
(end_heading_point[1] - center_y) / 2 + center_y)
heading_points = [start_heading_point, end_heading_point]
draw.line(heading_points, fill=color, width=thickness)
# Compute extremes so we can anchor the labels to them.
xs = [x[0] for x in points]
ys = [x[1] for x in points]
left = np.min(xs)
bottom = np.min(ys)
top = np.max(ys)
try:
font = ImageFont.truetype('arial.ttf', 24)
except IOError:
font = ImageFont.load_default()
text_width, text_height = font.getsize(display_str)
margin = np.ceil(0.05 * text_height)
if text_loc == 'TOP':
text_bottom = top
else:
text_bottom = bottom + 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)
np.copyto(original_image, np.array(image))
def TransformBBoxesToTopDown(bboxes, car_to_image_transform=None):
"""Convert bounding boxes from car coordinates to top down pixel coordinates.
Args:
bboxes: A (batch, nbboxes, 4 or 5) np.float32 tensor containing bounding box
xywhh in car coordinates (smooth adjusted by car pose).
car_to_image_transform: An optional Transform object. If None, this will be
created using _CarToImageTransform.
Returns:
np.array of shape (batch, nbboxes) containing the bounding boxes in top down
image space.
"""
if car_to_image_transform is None:
car_to_image_transform = _CarToImageTransform()
batch_size = np.shape(bboxes)[0]
nbboxes = np.shape(bboxes)[1]
transformed_boxes = np.zeros_like(bboxes)
for batch_id in range(batch_size):
for box_id in range(nbboxes):
# TODO(vrv): When we predict heading, we should assert
# that the length of bbox_data is 5.
bbox_data = np.squeeze(bboxes[batch_id, box_id, :])
x, y, width, height = (bbox_data[0], bbox_data[1], bbox_data[2],
bbox_data[3])
if len(bbox_data) == 5:
heading = bbox_data[4]
else:
heading = 0.0
# Skip boxes that cannot be visualized.
if width <= 0 or height <= 0:
continue
if any([np.isnan(c) or np.isinf(c) for c in bbox_data]):
continue
# Bounding boxes are in car coordinates (smooth adjusted by car pose).
# Transform from car coordinates to new coordinates.
bbox_car = transform_util.Box2D(x, y, width, height, heading)
bbox_transformed = bbox_car.Apply(car_to_image_transform)
bbox_values = bbox_transformed.AsNumpy()
if len(bbox_data) == 4:
bbox_values = bbox_values[:-1]
transformed_boxes[batch_id, box_id, :] = bbox_values
return transformed_boxes
def testBox2DTransform(self):
# Take the box from above and apply a car-image transform.
box = transform_util.Box2D(1.0, 1.0, 2., 1., 0.)
transform = transform_util.MakeCarToImageTransform(
pixels_per_meter=10.0, image_ref_x=250, image_ref_y=750, flip_axes=True)
new_box = box.Apply(transform)
# The center flips across the x=y axis to -1, -1. After the scaling and
# translation, the box should be centered (240, 740). Because the box flips
# across the axis, the width and length get flipped from 2, 1 to [10, 20].
#
# The flip axes should cause the heading to go from 0. to -pi/2.
self.assertAllClose([240., 740., 10., 20., -np.pi / 2.], new_box.AsNumpy())
# Check ymin/xmin/ymax/xmax: the rectangle is now longer in the y-dimension
# than the x-dimension.
self.assertAllClose((730., 235., 750., 245.), new_box.Extrema())
def DrawBBoxesOnImages(images, bboxes, box_weights, labels, class_id_to_name,
groundtruth):
"""Draw ground truth boxes on top down image.
Args:
images: A 4D uint8 array (batch, height, width, depth) of images to draw on
top of.
bboxes: A (batch, nbboxes, 4 or 5) np.float32 tensor containing bounding box
xywhh to draw specified in top down pixel values.
box_weights: A (batch, nbboxes) float matrix indicating the predicted score
of the box. If the score is 0.0, no box is drawn.
labels: A (batch, nbboxes) integer matrix indicating the true or predicted
label indices.
class_id_to_name: Dictionary mapping from class id to name.
groundtruth: Boolean indicating whether bounding boxes are ground truth.
Returns:
'images' with the bboxes drawn on top.
"""
# Assert 4d shape.
assert len(np.shape(images)) == 4
if np.shape(images)[3] == 1:
# Convert from grayscale to RGB.
images = np.tile(images, (1, 1, 1, 3))
# Assert channel dimension is 3 dimensional.
assert np.shape(images)[3] == 3
batch_size = np.shape(images)[0]
nbboxes = np.shape(bboxes)[1]
for batch_id in range(batch_size):
image = images[batch_id, :, :, :]
# Draw a box for each box and label if weights is 1.
transformed_boxes = []
box_scores = []
label_ids = []
for box_id in range(nbboxes):
box_weight = box_weights[batch_id, box_id]
# If there is no box to draw, continue.
if box_weight == 0.0:
continue
# TODO(vrv): When we predict heading, we should assert
# that the length of bbox_data is 5.
bbox_data = np.squeeze(bboxes[batch_id, box_id, :])
if len(bbox_data) == 5:
x, y, width, length, heading = bbox_data
else:
x, y, width, length = bbox_data
heading = 0.0
# Check whether we can draw the box.
bbox = transform_util.Box2D(x, y, width, length, heading)
ymin, xmin, ymax, xmax = bbox.Extrema()
if ymin == 0 and xmin == 0 and ymax == 0 and xmax == 0:
continue
# TODO(vrv): Support drawing boxes on the edge of the
# image.
if (xmin < 0 or ymin < 0 or xmax >= image.shape[1]
or ymax >= image.shape[0]):
continue
# We can draw a box on the image, so fill in the score, the boxes,
# and the label.
transformed_boxes.append([x, y, width, length, heading])
box_scores.append(box_weight)
label_ids.append(labels[batch_id, box_id])
if transformed_boxes:
transformed_boxes = np.stack(transformed_boxes, axis=0)
scores = None if groundtruth else np.array(box_scores)
text_loc = 'TOP' if groundtruth else 'BOTTOM'
VisualizeBoxes(image=image,
boxes=transformed_boxes,
classes=label_ids,
scores=scores,
class_id_to_name=class_id_to_name,
groundtruth_box_visualization_color='cyan',
skip_scores=groundtruth,
skip_labels=False,
text_loc=text_loc)
return images
def testBox2DCornerRotated(self):
# Like above but rotated 90 degrees; the order is important.
box = transform_util.Box2D(1.0, 1.0, 2., 1., np.pi / 2.)
self.assertAllClose(box.corners,
[[0., 0.5], [0., 1.5], [2., 1.5], [2., 0.5]])