vhad/body-pose-animation
1
0
Fork 0
mirror of https://github.com/gosticks/body-pose-animation.git synced 2025-10-16 11:45:42 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

add video renderer

Browse Source
This commit is contained in:
Wlad 2021-02-08 12:33:56 +01:00
parent 36f617899e
commit e9aeffb199
6 changed files with 164 additions and 51 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

27
example_fit.py
View File

@ -1,4 +1,5 @@
# library imports
from utils.render import make_video
import torch
import matplotlib.pyplot as plt
@ -21,7 +22,8 @@ init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_tra
model=model,
renderer=True,
dataset=dataset,
sample_index=sample_index
sample_index=sample_index,
offscreen=True
)
# configure PyTorch device and format
@ -45,7 +47,7 @@ camera.setup_visualization(r.init_keypoints, r.keypoints)
# train for pose
result, best, train_loss = train_pose_with_conf(
result, best, train_loss, step_imgs = train_pose_with_conf(
config=config,
model=model,
keypoints=keypoints,
@ -55,10 +57,17 @@ result, best, train_loss = train_pose_with_conf(
device=device,
)
fig, ax = plt.subplots()
name = getfilename_from_conf(config=config, index=sample_index)
ax.plot(train_loss[1::], label='sgd')
ax.set(xlabel="Training iteration", ylabel="Loss", title='Training loss')
fig.savefig("results/" + name + ".png")
ax.legend()
plt.show()
make_video(step_imgs, "test.avi")
# color = r.get_snapshot()
# plt.imshow(color)
# plt.show()
# fig, ax = plt.subplots()
# name = getfilename_from_conf(config=config, index=sample_index)
# ax.plot(train_loss[1::], label='sgd')
# ax.set(xlabel="Training iteration", ylabel="Loss", title='Training loss')
# fig.savefig("results/" + name + ".png")
# ax.legend()
# plt.show()

55
example_fit_anim.py
View File

@ -1,6 +1,8 @@
import pickle
import time
from utils.render import make_video
import torch
from tqdm.std import tqdm
from dataset import SMPLyDataset
from model import *
@ -17,6 +19,7 @@ FINISH_IDX = 200 # choose a big number to optimize for all frames in samples d
RUN_OPTIMIZATION = False
final_poses = [] # optimized poses array that is saved for playing the animation
result_image = []
idx = START_IDX
@ -27,15 +30,15 @@ def get_next_frame(idx):
:param idx: index of the frame
:return: tuple of keypoints, conf and image path
"""
keypoints = dataset[idx]
keypoints, keypoints_conf = dataset[idx]
if keypoints is None:
return
image_path = dataset.get_image_path(idx)
return keypoints[0], keypoints[1], image_path
return keypoints, keypoints_conf, image_path
device = torch.device('cpu')
dtype = torch.float
dtype = torch.float32
config = load_config()
dataset = SMPLyDataset.from_config(config)
@ -57,16 +60,15 @@ joints = model_out.joints.detach().cpu().numpy().squeeze()
Optimization part without visualization
'''
if RUN_OPTIMIZATION:
while get_next_frame(idx) is not None and idx <= FINISH_IDX:
keypoints, confidence, img_path = get_next_frame(idx)
for idx in dataset:
est_scale = estimate_scale(joints, keypoints)
# apply scaling to keypoints
keypoints = keypoints * est_scale
init_joints = get_torso(joints)
init_keypoints = get_torso(keypoints)
init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training(
model=model,
renderer=True,
offscreen=True,
dataset=dataset,
sample_index=idx
)
camera = TorchCameraEstimate(
model,
@ -86,9 +88,9 @@ if RUN_OPTIMIZATION:
config=config,
model=model,
keypoints=keypoints,
keypoint_conf=confidence,
keypoint_conf=conf,
camera=camera,
renderer=None,
renderer=r,
device=device
)
@ -142,6 +144,31 @@ def replay_animation(file, start_frame=0, end_frame=None, with_background=False,
time.sleep(1 / fps)
def video_from_pkl(filename, video_name):
with open(filename, "rb") as fp:
final_poses = pickle.load(fp)
save_to_video(final_poses, video_name)
def save_to_video(poses, video_name, fps=30):
r = DefaultRenderer(
offscreen=True
)
r.start()
model_anim = SMPLyModel.model_from_conf(config)
frames = []
for body_pose, cam_trans in tqdm(poses):
r.render_model_with_tfs(model_anim, body_pose, keep_pose=True,
render_joints=False, transforms=cam_trans)
frames.append(r.get_snapshot(False))
make_video(frames, video_name, fps)
'''
Play the animation.
'''

54
renderer.py
View File

@ -14,9 +14,12 @@ class Renderer:
camera_pose=None,
width=1920,
height=1080,
light_color=[0.3, 0.3, 0.3, 1.0]
light_color=[0.3, 0.3, 0.3, 1.0],
offscreen=False
) -> None:
super().__init__()
self.use_offscreen = offscreen
self.run_in_thread = False
self.scene = pyrender.Scene(
ambient_light=light_color
@ -50,12 +53,22 @@ class Renderer:
viewport_size = (self.width, self.height)
self.run_in_thread = run_in_thread
self.viewer = pyrender.Viewer(
self.scene,
run_in_thread=run_in_thread,
use_reymond_lighting=use_reymond_lighting,
viewport_size=tuple(d // 2 for d in viewport_size)
)
if self.use_offscreen:
self.offscreen = pyrender.OffscreenRenderer(
# self.scene,
# run_in_thread=run_in_thread,
# use_reymond_lighting=use_reymond_lighting,
viewport_width=(viewport_size[0]),
viewport_height=(viewport_size[1])
)
else:
self.viewer = pyrender.Viewer(
self.scene,
run_in_thread=run_in_thread,
use_reymond_lighting=use_reymond_lighting,
viewport_size=tuple(d // 2 for d in viewport_size)
)
def stop(self):
self.viewer.close_external()
@ -63,7 +76,7 @@ class Renderer:
pass
def requires_lock(self):
return self.run_in_thread and self.viewer
return not self.use_offscreen and self.run_in_thread and self.viewer
def release(self):
if self.requires_lock():
@ -212,8 +225,8 @@ class Renderer:
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
self.render_image(img, scale, name)
def render_image(self, image, scale, name=None):
_ = render_image_plane(self.scene, image, scale, name)
def render_image(self, image, scale, name="source_img"):
self.image = render_image_plane(self.scene, image, scale, name)
def set_homog_group_transform(self, group_name, rotation, translation):
# create pose matrix
@ -281,6 +294,27 @@ class Renderer:
if self.requires_lock():
self.viewer.render_lock.release()
def get_snapshot(self, show_image=False, camera_pose=None):
"""get snapshot of the current renderer, only works in offscreen mode
"""
if not self.use_offscreen:
print("[error] get_snapshot only works when used with offscreen renderer")
return None, None
if not show_image:
self.scene.remove_node(self.image)
color, depth = self.offscreen.render(
self.scene
)
# revert renderer changes
if not show_image:
self.scene.add_node(self.image)
return color
class DefaultRenderer(Renderer):
"""Utility class for easier default renderer setup

53
train_pose.py
View File

@ -16,32 +16,51 @@ from renderer import Renderer
def train_pose(
model: smplx.SMPL,
# current datapoints
keypoints,
keypoint_conf,
# 3D to 2D camera layer
camera: SimpleCamera,
# model type
model_type="smplx",
learning_rate=1e-3,
# pytorch config
device=torch.device('cuda'),
dtype=torch.float32,
renderer: Renderer = None,
# optimizer settings
optimizer=None,
optimizer_type="LBFGS",
learning_rate=1e-3,
iterations=60,
patience=10,
# configure loss function
useBodyPrior=False,
body_prior_weight=2,
useAnglePrior=False,
useConfWeights=False,
angle_prior_weight=0.5,
use_angle_sum_loss=False,
angle_sum_weight=0.1,
patience=10,
body_prior_weight=2,
angle_prior_weight=0.5,
body_mean_loss=False,
body_mean_weight=0.01
body_mean_weight=0.01,
useConfWeights=False,
# renderer options
renderer: Renderer = None,
render_steps=True,
render_offscreen=True
):
print("[pose] starting training")
print("[pose] dtype=", dtype)
offscreen_step_output = []
loss_layer = torch.nn.MSELoss().to(device=device, dtype=dtype) # MSELoss()
clip_loss_layer = AngleClipper().to(device=device, dtype=dtype)
@ -52,7 +71,7 @@ def train_pose(
# setup keypoint data
keypoints = torch.tensor(keypoints).to(device=device, dtype=dtype)
# get a list of openpose conf values
keypoints_conf = torch.tensor(keypoint_conf).to(device=device, dtype=dtype)
# keypoints_conf = torch.tensor(keypoint_conf).to(device=device, dtype=dtype)
# create filter layer to ignore unused joints, keypoints during optimization
filter_layer = JointFilter(
@ -89,7 +108,7 @@ def train_pose(
pbar = tqdm(total=iterations)
def predict():
pose_extra = None
# pose_extra = None
# if useBodyPrior:
# body = vposer_layer()
@ -120,7 +139,7 @@ def train_pose(
body_prior_loss = 0.0
if useBodyPrior:
# apply pose prior loss.
body_prior_loss = latent_body.pow(
body_prior_loss = latent_pose.pow(
2).sum() * body_prior_weight
angle_prior_loss = 0.0
@ -162,9 +181,6 @@ def train_pose(
pred = predict()
loss = optim_closure()
# if t % 5 == 0:
# time.sleep(5)
# compute loss
cur_loss = loss.item()
@ -184,15 +200,18 @@ def train_pose(
pbar.set_description("Error %f" % cur_loss)
pbar.update(1)
if renderer is not None:
if renderer is not None and render_steps:
R = camera.trans.detach().cpu().numpy().squeeze()
renderer.render_model_with_tfs(
model, pose_layer.cur_out, keep_pose=True, transforms=R)
if render_offscreen:
offscreen_step_output.append(renderer.get_snapshot())
# renderer.set_group_pose("body", R)
pbar.close()
print("Final result:", loss.item())
return pose_layer.cur_out, best_pose, loss_history
return pose_layer.cur_out, best_pose, loss_history, offscreen_step_output
def train_pose_with_conf(
@ -204,6 +223,7 @@ def train_pose_with_conf(
device=torch.device('cpu'),
dtype=torch.float32,
renderer: Renderer = None,
render_steps=True
):
# configure PyTorch device and format
@ -244,5 +264,6 @@ def train_pose_with_conf(
body_mean_loss=config['pose']['bodyMeanLoss']['enabled'],
body_mean_weight=config['pose']['bodyMeanLoss']['weight'],
use_angle_sum_loss=config['pose']['angleSumLoss']['enabled'],
angle_sum_weight=config['pose']['angleSumLoss']['weight']
angle_sum_weight=config['pose']['angleSumLoss']['weight'],
render_steps=render_steps
)

4
utils/general.py
View File

@ -152,7 +152,7 @@ def get_new_filename():
return result_prefix + str(num + 1) + ".pkl"
def setup_training(model, dataset, sample_index, renderer=True):
def setup_training(model, dataset, sample_index, renderer=True, offscreen=False):
keypoints, conf = dataset[sample_index]
img_path = dataset.get_image_path(sample_index)
@ -179,7 +179,7 @@ def setup_training(model, dataset, sample_index, renderer=True):
if renderer:
# setup renderer
r = DefaultRenderer()
r = DefaultRenderer(offscreen=offscreen)
r.setup(
model=model,
model_out=model_out,

22
utils/render.py
View File

@ -2,6 +2,28 @@ from typing import List, Set, Dict, Tuple, Optional
import numpy as np
import trimesh
import pyrender
import cv2
from tqdm import tqdm
def make_video(images, video_name: str, fps=5):
images = np.array(images)
print(images.shape)
width = images.shape[2]
height = images.shape[1]
video = cv2.VideoWriter(
video_name, 0, fps, (width, height), True)
print("creating video with size", width, height)
for idx in tqdm(range(len(images))):
img = images[idx]
im_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
video.write(im_rgb)
video.release()
def render_model(