mirror of
https://github.com/gosticks/body-pose-animation.git
synced 2025-10-16 11:45:42 +00:00
197 lines
6.3 KiB
Python
197 lines
6.3 KiB
Python
from dataset import SMPLyDataset
|
|
import pickle
|
|
from typing import Tuple
|
|
from model import SMPLyModel
|
|
from renderer import DefaultRenderer
|
|
import cv2
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
from scipy import interpolate
|
|
|
|
|
|
def make_video(images, video_name: str, fps=30, ext: str = "mp4", post_process_frame=None):
|
|
images = np.array(images)
|
|
width = images.shape[2]
|
|
height = images.shape[1]
|
|
|
|
fourcc = 0
|
|
if ext == "mp4":
|
|
fourcc = cv2.VideoWriter_fourcc(*'MP4V')
|
|
|
|
video_name = video_name + "." + ext
|
|
|
|
video = cv2.VideoWriter(
|
|
video_name, fourcc, fps, (width, height), True)
|
|
|
|
for idx in tqdm(range(len(images))):
|
|
img = images[idx]
|
|
im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
|
|
if post_process_frame is not None:
|
|
img_rgb = post_process_frame(img=im_rgb, idx=idx)
|
|
|
|
video.write(im_rgb)
|
|
|
|
video.release()
|
|
print("video saved to:", video_name)
|
|
|
|
|
|
def video_from_pkl(filename, video_name, config, ext: str = "mp4"):
|
|
with open(filename, "rb") as fp:
|
|
model_outs = pickle.load(fp)
|
|
save_to_video(model_outs, video_name, config)
|
|
|
|
|
|
def save_to_video(
|
|
sample_output: Tuple,
|
|
video_name: str,
|
|
config: object,
|
|
fps=30,
|
|
include_thumbnail=True,
|
|
thumbnail_size=0.2,
|
|
start_frame_offset=0,
|
|
dataset: SMPLyDataset = None,
|
|
interpolation_target=None
|
|
):
|
|
"""
|
|
Renders a video from pose, camera tuples. Additionally interpolation can be used to smooth out the animation
|
|
|
|
Args:
|
|
sample_output (Tuple): A tuple of body pose vertices and a camera transformation
|
|
video_name (str): name for the resulting video file (can also be a path)
|
|
config (object): general run config
|
|
fps (int, optional): animation base fps. Defaults to 30.
|
|
interpolation_target (int, optional): expand animation fps via interpolation to this target. Defaults to 60.
|
|
"""
|
|
r = DefaultRenderer(
|
|
offscreen=True
|
|
)
|
|
r.start()
|
|
|
|
model_anim = SMPLyModel.model_from_conf(config)
|
|
|
|
if interpolation_target is not None:
|
|
if interpolation_target % fps != 0:
|
|
print("[error] interpolation target must be a multiple of fps")
|
|
return
|
|
inter_ratio = int(interpolation_target / fps)
|
|
num_intermediate = inter_ratio - 1
|
|
sample_output = interpolate_poses(sample_output, num_intermediate)
|
|
else:
|
|
sample_output = [
|
|
(
|
|
out.vertices.detach().cpu().numpy()[0],
|
|
cam
|
|
) for out, cam in sample_output]
|
|
frames = []
|
|
print("[export] rendering animation frames...", sample_output[0][0].shape)
|
|
|
|
# just use the first transform
|
|
cam_transform = sample_output[0][1]
|
|
|
|
for vertices, cam_trans in tqdm(sample_output):
|
|
r.render_model_geometry(
|
|
faces=model_anim.faces,
|
|
vertices=vertices,
|
|
pose=cam_trans # cam_transform,
|
|
)
|
|
frames.append(r.get_snapshot())
|
|
|
|
target_fps = fps
|
|
if interpolation_target is not None:
|
|
target_fps = interpolation_target
|
|
|
|
def post_process_frame(img, idx: int):
|
|
if not include_thumbnail:
|
|
return img
|
|
# account for start from frames not zero
|
|
idx = start_frame_offset + idx
|
|
frame_idx = idx
|
|
if interpolation_target is not None:
|
|
# account for possible interpolation
|
|
frame_idx = int(idx / inter_ratio)
|
|
img_path = dataset.get_image_path(frame_idx)
|
|
overlay = cv2.imread(img_path)
|
|
|
|
if overlay is None:
|
|
print("[error] image could not be ", img_path)
|
|
return img
|
|
|
|
overlay = cv2.resize(
|
|
overlay,
|
|
dsize=(
|
|
int(overlay.shape[1] * thumbnail_size),
|
|
int(overlay.shape[0] * thumbnail_size)
|
|
))
|
|
img[0:overlay.shape[0], 0:overlay.shape[1]] = overlay
|
|
return img
|
|
|
|
make_video(frames, video_name, target_fps,
|
|
post_process_frame=post_process_frame)
|
|
|
|
|
|
def make_video_with_pip(frames, pip_image_path, video_name: str, fps=30, ext: str = "mp4", image_size=0.2):
|
|
"""renders a video with a pip frame in the corner
|
|
"""
|
|
|
|
def post_process_frame(img, idx: int):
|
|
overlay = cv2.imread(pip_image_path)
|
|
|
|
if overlay is None:
|
|
print("[error] image could not be ", pip_image_path)
|
|
return img
|
|
|
|
overlay = cv2.resize(
|
|
overlay,
|
|
dsize=(
|
|
int(overlay.shape[1] * image_size),
|
|
int(overlay.shape[0] * image_size)
|
|
))
|
|
img[0:overlay.shape[0], 0:overlay.shape[1]] = overlay
|
|
return img
|
|
|
|
make_video(frames, video_name, fps,
|
|
post_process_frame=post_process_frame)
|
|
|
|
|
|
def interpolate_poses(poses, num_intermediate=5):
|
|
"""
|
|
Interpolate vertices and cameras between pairs of frames by adding intermediate results
|
|
|
|
:param poses: optimized poses
|
|
:param num_intermediate: amount of intermediate results to insert between each pair of frames
|
|
:return: interpolated poses, list of tuples (body_pose, camera_pose)
|
|
"""
|
|
new_poses = []
|
|
for i in range(len(poses) - 1):
|
|
if len(poses) < 2:
|
|
return poses
|
|
else:
|
|
# Shape of one matrix of vertices = torch.Size([1, 10475, 3])
|
|
pose_1 = poses[i][0].vertices.detach().cpu().numpy()
|
|
pose_2 = poses[i + 1][0].vertices.detach().cpu().numpy()
|
|
poses_pair = np.concatenate((pose_1, pose_2), axis=0)
|
|
|
|
camera_1 = np.expand_dims(poses[i][1], axis=0)
|
|
camera_2 = np.expand_dims(poses[i + 1][1], axis=0)
|
|
camera_pair = np.concatenate((camera_1, camera_2), axis=0)
|
|
|
|
x = np.arange(poses_pair.shape[0])
|
|
f1 = interpolate.interp1d(x, poses_pair, axis=0)
|
|
f2 = interpolate.interp1d(x, camera_pair, axis=0)
|
|
|
|
evenly_spaced_points = np.linspace(
|
|
x[0], x[-1], (poses_pair.shape[0] - 1) * (num_intermediate + 1) + 1)
|
|
|
|
new_frames = f1(evenly_spaced_points)
|
|
new_cameras = f2(evenly_spaced_points)
|
|
|
|
arr = [(new_frames[i], new_cameras[i])
|
|
for i in range(new_frames.shape[0])]
|
|
if 0 < i < len(poses) - 1:
|
|
# remove first frame that was already added in the last interpolation
|
|
arr.pop(0)
|
|
new_poses += arr
|
|
|
|
return new_poses
|