File size: 3,637 Bytes
f1586f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
"""Visualization utility functions."""
import colorsys
import random
from typing import List, Optional, Sequence, Tuple
import numpy as np
# Generate random colormaps for visualizing different points.
def get_colors(num_colors: int) -> List[Tuple[int, int, int]]:
"""Gets colormap for points."""
colors = []
for i in np.arange(0.0, 360.0, 360.0 / num_colors):
hue = i / 360.0
lightness = (50 + np.random.rand() * 10) / 100.0
saturation = (90 + np.random.rand() * 10) / 100.0
color = colorsys.hls_to_rgb(hue, lightness, saturation)
colors.append(
(int(color[0] * 255), int(color[1] * 255), int(color[2] * 255))
)
random.shuffle(colors)
return colors
def paint_point_track(
frames: np.ndarray,
point_tracks: np.ndarray,
visibles: np.ndarray,
colormap: Optional[List[Tuple[int, int, int]]] = None,
) -> np.ndarray:
"""Converts a sequence of points to color code video.
Args:
frames: [num_frames, height, width, 3], np.uint8, [0, 255]
point_tracks: [num_points, num_frames, 2], np.float32, [0, width / height]
visibles: [num_points, num_frames], bool
colormap: colormap for points, each point has a different RGB color.
Returns:
video: [num_frames, height, width, 3], np.uint8, [0, 255]
"""
num_points, num_frames = point_tracks.shape[0:2]
if colormap is None:
colormap = get_colors(num_colors=num_points)
height, width = frames.shape[1:3]
dot_size_as_fraction_of_min_edge = 0.015
radius = int(round(min(height, width) * dot_size_as_fraction_of_min_edge))
diam = radius * 2 + 1
quadratic_y = np.square(np.arange(diam)[:, np.newaxis] - radius - 1)
quadratic_x = np.square(np.arange(diam)[np.newaxis, :] - radius - 1)
icon = (quadratic_y + quadratic_x) - (radius**2) / 2.0
sharpness = 0.15
icon = np.clip(icon / (radius * 2 * sharpness), 0, 1)
icon = 1 - icon[:, :, np.newaxis]
icon1 = np.pad(icon, [(0, 1), (0, 1), (0, 0)])
icon2 = np.pad(icon, [(1, 0), (0, 1), (0, 0)])
icon3 = np.pad(icon, [(0, 1), (1, 0), (0, 0)])
icon4 = np.pad(icon, [(1, 0), (1, 0), (0, 0)])
video = frames.copy()
for t in range(num_frames):
# Pad so that points that extend outside the image frame don't crash us
image = np.pad(
video[t],
[
(radius + 1, radius + 1),
(radius + 1, radius + 1),
(0, 0),
],
)
for i in range(num_points):
# The icon is centered at the center of a pixel, but the input coordinates
# are raster coordinates. Therefore, to render a point at (1,1) (which
# lies on the corner between four pixels), we need 1/4 of the icon placed
# centered on the 0'th row, 0'th column, etc. We need to subtract
# 0.5 to make the fractional position come out right.
x, y = point_tracks[i, t, :] + 0.5
x = min(max(x, 0.0), width)
y = min(max(y, 0.0), height)
if visibles[i, t]:
x1, y1 = np.floor(x).astype(np.int32), np.floor(y).astype(np.int32)
x2, y2 = x1 + 1, y1 + 1
# bilinear interpolation
patch = (
icon1 * (x2 - x) * (y2 - y)
+ icon2 * (x2 - x) * (y - y1)
+ icon3 * (x - x1) * (y2 - y)
+ icon4 * (x - x1) * (y - y1)
)
x_ub = x1 + 2 * radius + 2
y_ub = y1 + 2 * radius + 2
image[y1:y_ub, x1:x_ub, :] = (1 - patch) * image[
y1:y_ub, x1:x_ub, :
] + patch * np.array(colormap[i])[np.newaxis, np.newaxis, :]
# Remove the pad
video[t] = image[
radius + 1 : -radius - 1, radius + 1 : -radius - 1
].astype(np.uint8)
return video
|