mirror of
https://github.com/gosticks/body-pose-animation.git
synced 2025-10-16 11:45:42 +00:00
102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
# library imports
|
|
import pickle
|
|
import torch
|
|
from utils.video import make_video
|
|
from tqdm.auto import trange
|
|
|
|
# local imports
|
|
from train_pose import train_pose_with_conf
|
|
from model import SMPLyModel
|
|
from utils.general import get_new_filename, setup_training
|
|
from camera_estimation import TorchCameraEstimate
|
|
|
|
|
|
def optimize_sample(sample_index, dataset, config, device=torch.device('cpu'), dtype=torch.float32, offscreen=False, verbose=False, display_result=False):
|
|
# prepare data and SMPL model
|
|
model = SMPLyModel.model_from_conf(config)
|
|
init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training(
|
|
model=model,
|
|
renderer=True,
|
|
dataset=dataset,
|
|
sample_index=sample_index,
|
|
offscreen=offscreen
|
|
)
|
|
|
|
camera = TorchCameraEstimate(
|
|
model,
|
|
keypoints=keypoints,
|
|
renderer=r,
|
|
device=device,
|
|
dtype=dtype,
|
|
image_path=img_path,
|
|
est_scale=est_scale,
|
|
verbose=verbose,
|
|
use_progress_bar=verbose
|
|
)
|
|
|
|
camera_transformation, camera_int, camera_params = camera.get_results()
|
|
|
|
if not offscreen:
|
|
# render camera to the scene
|
|
camera.setup_visualization(r.init_keypoints, r.keypoints)
|
|
|
|
# train for pose
|
|
pose, loss_history, step_imgs = train_pose_with_conf(
|
|
config=config,
|
|
model=model,
|
|
keypoints=keypoints,
|
|
keypoint_conf=conf,
|
|
camera=camera,
|
|
renderer=r,
|
|
device=device,
|
|
use_progress_bar=verbose
|
|
)
|
|
|
|
if display_result:
|
|
r.wait_for_close()
|
|
|
|
return pose, camera_transformation, loss_history, step_imgs
|
|
|
|
|
|
def create_animation(dataset, config, start_idx=0, end_idx=None, device=torch.device('cpu'), dtype=torch.float32, offscreen=False, verbose=False, save_to_file=False):
|
|
final_poses = []
|
|
|
|
if end_idx is None:
|
|
end_idx = len(dataset)
|
|
|
|
for idx in trange(end_idx - start_idx, desc='Optimizing'):
|
|
idx = start_idx + idx
|
|
|
|
final_pose, cam_trans, train_loss, step_imgs = optimize_sample(
|
|
idx,
|
|
dataset,
|
|
config,
|
|
offscreen=True
|
|
)
|
|
|
|
if verbose:
|
|
print("Optimization of", idx, "frames finished")
|
|
|
|
# print("\nPose optimization of frame", idx, "is finished.")
|
|
R = cam_trans.numpy().squeeze()
|
|
idx += 1
|
|
|
|
# append optimized pose and camera transformation to the array
|
|
final_poses.append((final_pose, R))
|
|
|
|
filename = None
|
|
|
|
if save_to_file:
|
|
'''
|
|
Save final_poses array into results folder as a pickle dump
|
|
'''
|
|
results_dir = config['output']['rootDir']
|
|
result_prefix = config['output']['prefix']
|
|
filename = results_dir + get_new_filename()
|
|
print("Saving results to", filename)
|
|
with open(filename, "wb") as fp:
|
|
pickle.dump(final_poses, fp)
|
|
print("Results have been saved to", filename)
|
|
|
|
return final_poses, filename
|