mirror of
https://github.com/gosticks/body-pose-animation.git
synced 2025-10-16 11:45:42 +00:00
52 lines
1.5 KiB
Python
52 lines
1.5 KiB
Python
import os
|
|
from utils.video import make_video_with_pip
|
|
from utils.graphs import render_loss_graph
|
|
from train import optimize_sample
|
|
|
|
# local imports
|
|
from utils.general import get_output_path_from_conf, load_config
|
|
from dataset import SMPLyDataset
|
|
|
|
# load and select sample
|
|
config = load_config()
|
|
dataset = SMPLyDataset.from_config(config=config)
|
|
sample_index = 55
|
|
save_to_video = True
|
|
|
|
if os.getenv('SAMPLE_INDEX') is not None:
|
|
sample_index = int(os.getenv('SAMPLE_INDEX'))
|
|
|
|
# train for pose
|
|
best_pose, camera_transformation, loss_history, step_imgs, loss_components = optimize_sample(
|
|
sample_index - 1,
|
|
dataset,
|
|
config,
|
|
interactive=False
|
|
)
|
|
|
|
config['pose']['lr'] = config['pose']['temporal']['lr']
|
|
config['pose']['iterations'] = config['pose']['temporal']['iterations']
|
|
|
|
# reuse model to train new sample
|
|
pose_temp, camera_transformation, loss_history, step_imgs, loss_components = optimize_sample(
|
|
sample_index,
|
|
dataset,
|
|
config,
|
|
interactive=False,
|
|
offscreen=True,
|
|
initial_pose=best_pose.body_pose.detach().clone().cpu(),
|
|
initial_orient=best_pose.global_orient.detach().clone().cpu()
|
|
)
|
|
|
|
if save_to_video:
|
|
img_path = dataset.get_image_path(sample_index)
|
|
make_video_with_pip(step_imgs, pip_image_path=img_path,
|
|
video_name="example_fit_temporal")
|
|
|
|
filename = get_output_path_from_conf(config) + ".png"
|
|
render_loss_graph(
|
|
loss_history=loss_history,
|
|
loss_components=loss_components,
|
|
save=True,
|
|
filename=filename)
|