146 lines
5.2 KiB
Python
146 lines
5.2 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch.nn.functional as F
|
|
from skimage.draw import circle_perimeter as circle
|
|
|
|
|
|
class Visualizer:
|
|
def __init__(
|
|
self,
|
|
kp_size=5,
|
|
draw_border=False,
|
|
colormap="gist_rainbow",
|
|
region_bg_color=(0, 0, 0),
|
|
):
|
|
self.kp_size = kp_size
|
|
self.draw_border = draw_border
|
|
self.colormap = plt.get_cmap(colormap)
|
|
self.region_bg_color = np.array(region_bg_color)
|
|
|
|
def draw_image_with_kp(self, image, kp_array):
|
|
image = np.copy(image)
|
|
spatial_size = np.array(image.shape[:2][::-1])[np.newaxis]
|
|
kp_array = spatial_size * (kp_array + 1) / 2
|
|
num_regions = kp_array.shape[0]
|
|
for kp_ind, kp in enumerate(kp_array):
|
|
rr, cc = circle(int(kp[1]), int(kp[0]), self.kp_size, shape=image.shape[:2])
|
|
image[rr, cc] = np.array(self.colormap(kp_ind / num_regions))[:3]
|
|
return image
|
|
|
|
def create_image_column_with_kp(self, images, kp):
|
|
image_array = np.array(
|
|
[self.draw_image_with_kp(v, k) for v, k in zip(images, kp)]
|
|
)
|
|
return self.create_image_column(image_array)
|
|
|
|
def create_image_column(self, images):
|
|
if self.draw_border:
|
|
images = np.copy(images)
|
|
images[:, :, [0, -1]] = (1, 1, 1)
|
|
images[:, :, [0, -1]] = (1, 1, 1)
|
|
return np.concatenate(list(images), axis=0)
|
|
|
|
def create_image_grid(self, *args):
|
|
out = []
|
|
for arg in args:
|
|
if type(arg) == tuple:
|
|
out.append(self.create_image_column_with_kp(arg[0], arg[1]))
|
|
else:
|
|
out.append(self.create_image_column(arg))
|
|
return np.concatenate(out, axis=1)
|
|
|
|
def visualize(self, driving, source, out):
|
|
images = []
|
|
|
|
# Source image with region centers
|
|
source = source.data.cpu()
|
|
source_region_params = out["source_region_params"]["shift"].data.cpu().numpy()
|
|
source = np.transpose(source, [0, 2, 3, 1])
|
|
images.append((source, source_region_params))
|
|
|
|
# Equivariance visualization
|
|
if "transformed_frame" in out:
|
|
transformed = out["transformed_frame"].data.cpu().numpy()
|
|
transformed = np.transpose(transformed, [0, 2, 3, 1])
|
|
transformed_kp = (
|
|
out["transformed_region_params"]["shift"].data.cpu().numpy()
|
|
)
|
|
images.append((transformed, transformed_kp))
|
|
|
|
# Driving image with region centers
|
|
driving_region_params = out["driving_region_params"]["shift"].data.cpu().numpy()
|
|
driving = driving.data.cpu().numpy()
|
|
driving = np.transpose(driving, [0, 2, 3, 1])
|
|
images.append((driving, driving_region_params))
|
|
|
|
# Deformed image
|
|
if "deformed" in out:
|
|
deformed = out["deformed"].data.cpu().numpy()
|
|
deformed = np.transpose(deformed, [0, 2, 3, 1])
|
|
images.append(deformed)
|
|
|
|
# Result
|
|
prediction = out["prediction"].data.cpu().numpy()
|
|
prediction = np.transpose(prediction, [0, 2, 3, 1])
|
|
images.append(prediction)
|
|
|
|
# Occlusion map
|
|
if "occlusion_map" in out:
|
|
occlusion_map = out["occlusion_map"].data.cpu().repeat(1, 3, 1, 1)
|
|
occlusion_map = F.interpolate(occlusion_map, size=source.shape[1:3]).numpy()
|
|
occlusion_map = np.transpose(occlusion_map, [0, 2, 3, 1])
|
|
images.append(occlusion_map)
|
|
|
|
# Heatmaps visualizations
|
|
if "heatmap" in out["driving_region_params"]:
|
|
driving_heatmap = F.interpolate(
|
|
out["driving_region_params"]["heatmap"], size=source.shape[1:3]
|
|
)
|
|
driving_heatmap = np.transpose(
|
|
driving_heatmap.data.cpu().numpy(), [0, 2, 3, 1]
|
|
)
|
|
images.append(
|
|
draw_colored_heatmap(
|
|
driving_heatmap, self.colormap, self.region_bg_color
|
|
)
|
|
)
|
|
|
|
if "heatmap" in out["source_region_params"]:
|
|
source_heatmap = F.interpolate(
|
|
out["source_region_params"]["heatmap"], size=source.shape[1:3]
|
|
)
|
|
source_heatmap = np.transpose(
|
|
source_heatmap.data.cpu().numpy(), [0, 2, 3, 1]
|
|
)
|
|
images.append(
|
|
draw_colored_heatmap(
|
|
source_heatmap, self.colormap, self.region_bg_color
|
|
)
|
|
)
|
|
|
|
image = self.create_image_grid(*images)
|
|
image = (255 * image).astype(np.uint8)
|
|
return image
|
|
|
|
|
|
def draw_colored_heatmap(heatmap, colormap, bg_color):
|
|
parts = []
|
|
weights = []
|
|
bg_color = np.array(bg_color).reshape((1, 1, 1, 3))
|
|
num_regions = heatmap.shape[-1]
|
|
for i in range(num_regions):
|
|
color = np.array(colormap(i / num_regions))[:3]
|
|
color = color.reshape((1, 1, 1, 3))
|
|
part = heatmap[:, :, :, i : (i + 1)]
|
|
part = part / np.max(part, axis=(1, 2), keepdims=True)
|
|
weights.append(part)
|
|
|
|
color_part = part * color
|
|
parts.append(color_part)
|
|
|
|
weight = sum(weights)
|
|
bg_weight = 1 - np.minimum(1, weight)
|
|
weight = np.maximum(1, weight)
|
|
result = sum(parts) / weight + bg_weight * bg_color
|
|
return result
|