47 lines
1.4 KiB
Python
47 lines
1.4 KiB
Python
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
from skelda import utils_view
|
|
|
|
# ==================================================================================================
|
|
|
|
|
|
def show_poses2d(bodies, images, joint_names, title=""):
|
|
num_imgs = len(images)
|
|
rowbreak = int(num_imgs / 2.0 + 0.5)
|
|
fig, axs = plt.subplots(2, rowbreak, figsize=(30, 20))
|
|
fig.suptitle(title, fontsize=20)
|
|
|
|
if isinstance(bodies, np.ndarray):
|
|
bodies = bodies.tolist()
|
|
|
|
# Draw skeletons into images
|
|
for i, image in enumerate(images):
|
|
colors = plt.cm.hsv(np.linspace(0, 1, len(bodies[i]), endpoint=False)).tolist()
|
|
colors = [[int(c[0] * 255), int(c[1] * 255), int(c[2] * 255)] for c in colors]
|
|
|
|
for j, body in enumerate(bodies[i]):
|
|
image = utils_view.draw_body_in_image(image, body, joint_names, colors[j])
|
|
|
|
# Rescale image range for plotting
|
|
images = [img / 255.0 for img in images]
|
|
|
|
if rowbreak == 1:
|
|
axs[0].imshow(images[0])
|
|
if len(images) == 2:
|
|
axs[1].imshow(images[1])
|
|
else:
|
|
# Optionally delete last empty plot
|
|
fig.delaxes(axs[1])
|
|
|
|
else:
|
|
for i in range(rowbreak):
|
|
axs[0][i].imshow(images[i])
|
|
if i + rowbreak < num_imgs:
|
|
axs[1][i].imshow(images[i + rowbreak])
|
|
else:
|
|
# Optionally delete last empty plot
|
|
fig.delaxes(axs[1][i])
|
|
|
|
return fig
|