update:

- update renderer (unify methods and simplify)
- expose more configs to the yaml file
- add interpolation to video
- fix unstable camera while in video mode

config.yaml

@@ -22,17 +22,23 @@ pose:
   device: cuda
   lr: 0.01
   optimizer: Adam # currently supported Adam, LBFGS
-  iterations: 150
+  iterations: 100
   useCameraIntrinsics: true
+  useOpenPoseConf: true # use openpose confidence to weight L2 distance loss
   bodyMeanLoss:
     enabled: false
     weight: 0.1
   bodyPrior:
     enabled: true
-    weight: 0.01
+    weight: 0.1
   anglePrior:
     enabled: true
-    weight: 0.01
+    weight: 0.05
+    # optional per joint configurations
+    angleIdx: [56, 53, 12, 9, 37, 40]
+    directions: [-1, 1, -1, -1, -1, -1]
+    # weights per joint
+    weights: [0.8, 0.8, 0.8, 0.8, 0.1, 0.1]
   angleLimitLoss:
     enabled: true
     weight: 0.1
@@ -41,7 +47,7 @@ pose:
     weight: 0.01
   intersectLoss:
     enabled: true
-    weight: 2.0
+    weight: 1.0
     maxCollisions: 8
     sigma: 0.5
   changeLoss:

example_fit_anim.py

@@ -1,7 +1,6 @@
 import pickle
 import time
 from train import create_animation
-from utils.video import video_from_pkl
 from dataset import SMPLyDataset
 from model import *
 from utils.general import *
@@ -9,7 +8,7 @@ from renderer import *
 from utils.general import rename_files, get_new_filename
 
 START_IDX = 1  # starting index of the frame to optimize for
-FINISH_IDX = 100   # choose a big number to optimize for all frames in samples directory
+FINISH_IDX = 300   # choose a big number to optimize for all frames in samples directory
 # if False, only run already saved animation without optimization
 RUN_OPTIMIZATION = True
 
@@ -44,8 +43,6 @@ if RUN_OPTIMIZATION:
         interpolate=False
     )
 
-# TODO: use current body pose and camera transform for next optimization?
-
 
 def replay_animation(file, start_frame=0, end_frame=None, with_background=False, fps=30, interpolated=False):
     r = Renderer()
@@ -54,25 +51,23 @@ def replay_animation(file, start_frame=0, end_frame=None, with_background=False,
     model_anim = SMPLyModel.model_from_conf(config)
 
     with open(file, "rb") as fp:
-        model_outs = pickle.load(fp)
+        results = pickle.load(fp)
 
     if end_frame is None:
-        end_frame = len(model_outs)
+        end_frame = len(results)
 
-    for i in range(start_frame, end_frame):
-        body_pose = model_outs[i][0]
-        camera_transform = model_outs[i][1]
+    for model, camera_transform in results[start_frame::]:
+        if interpolated:
+            vertices = model
+        else:
+            vertices = model.vertices
 
-        if with_background:
-            # Changing image is too jerky, because the image has to be removed and added each time
-            pass
-            # img_path = samples_dir + "/" + str(i) + ".png"
-            # if r.get_node("image") is not None:
-            #     r.remove_node("image")
-            # r.render_image_from_path(img_path, name="image", scale=est_scale)
+        r.render_model_geometry(
+            faces=model_anim.faces,
+            vertices=vertices,
+            pose=camera_transform
+        )
 
-        r.render_model_with_tfs(model_anim, body_pose, keep_pose=True,
-                                render_joints=False, transforms=camera_transform, interpolated=interpolated)
         time.sleep(1 / fps)
 
 
@@ -86,8 +81,4 @@ else:
     result_prefix = config['output']['prefix']
     anim_file = results_dir + result_prefix + "0.pkl"
 
-video_name = getfilename_from_conf(
-    config) + "-" + str(START_IDX) + "-" + str(FINISH_IDX)
-
-video_from_pkl(anim_file, video_name, config)
-# replay_animation(anim_file, interpolated=True)
+replay_animation(anim_file, interpolated=True)

example_render_video.py

@@ -0,0 +1,35 @@
+import pickle
+import time
+from train import create_animation
+from utils.video import save_to_video
+from dataset import SMPLyDataset
+from model import *
+from utils.general import *
+from renderer import *
+from utils.general import rename_files, get_new_filename
+
+START_IDX = 1  # starting index of the frame to optimize for
+FINISH_IDX = 10   # choose a big number to optimize for all frames in samples directory
+
+result_image = []
+idx = START_IDX
+
+config = load_config()
+dataset = SMPLyDataset.from_config(config)
+model = SMPLyModel.model_from_conf(config)
+
+model_outs, filename = create_animation(
+    dataset,
+    config,
+    START_IDX,
+    FINISH_IDX,
+    verbose=False,
+    offscreen=True,
+    save_to_file=False,
+    interpolate=False
+)
+
+video_name = getfilename_from_conf(
+    config) + "-" + str(START_IDX) + "-" + str(FINISH_IDX)
+
+save_to_video(model_outs, video_name, config, interpolation_target=120)

modules/angle_prior.py

@@ -9,8 +9,8 @@ class AnglePriorsLoss(nn.Module):
         dtype=torch.float32,
         angle_idx=[56, 53, 12, 9, 37, 40],
         directions=[1, 1, -1, -1, 1, -1],
-        global_weight=1,
-        weights=[0.4, 0.4, 0.3, 0.3, 0.01, 0.01]
+        weight=1,
+        weights=[0.8, 0.8, 0.7, 0.7, 0.3, 0.3]
     ):
         super(AnglePriorsLoss, self).__init__()
 
@@ -37,7 +37,7 @@ class AnglePriorsLoss(nn.Module):
 
         self.register_buffer(
             "global_weight",
-            torch.tensor(global_weight, dtype=dtype).to(device=device)
+            torch.tensor(weight, dtype=dtype).to(device=device)
         )
 
     def forward(self, pose, joints, points, keypoints, raw_output):

modules/utils.py

@@ -59,10 +59,12 @@ def get_loss_layers(config, model: smplx.SMPL, device, dtype):
             weight=config['pose']['bodyPrior']['weight']))
 
     if config['pose']['anglePrior']['enabled']:
+        params = get_layer_config(config, "anglePrior")
         extra_loss_layers.append(AnglePriorsLoss(
             device=device,
-            global_weight=config['pose']['anglePrior']['weight'],
-            dtype=dtype))
+            dtype=dtype,
+            **params
+        ))
 
     if config['pose']['angleSumLoss']['enabled']:
         extra_loss_layers.append(AngleSumLoss(

renderer.py

@@ -3,7 +3,7 @@ import numpy as np
 from pyrender import scene
 from smplx import SMPLLayer
 from smplx.body_models import SMPL
-from utils.render import render_model, render_model_with_tfs, render_points, render_camera, render_image_plane
+from utils.render import render_model, render_model_geometry, render_points, render_camera, render_image_plane
 import pyrender
 
 
@@ -163,63 +163,52 @@ class Renderer:
             self,
             model: SMPLLayer,
             model_out,
-            color=[1.0, 0.3, 0.3, 0.8],
             replace=True,
-            keep_pose=True,
-            render_joints=True
+            transform=None,
+            render_joints=True,
+            **kwargs
+
     ):
         if model_out is None:
             model_out = model()
 
-        if keep_pose:
-            node = self.get_node("body_mesh")
-            # if node is not None:
-            #original_pose = node.pose
+        joints = None
 
         if render_joints:
-            self.render_joints(
-                model_out.joints.detach().cpu().numpy().squeeze())
+            joints = model_out.joints.detach().cpu().numpy().squeeze()
 
-        self.remove_from_group("body", "body_mesh")
+        return self.render_model_geometry(
+            replace=replace,
+            pose=transform,
+            faces=model.faces,
+            joints=joints,
+            vertices=model_out.vertices.detach().cpu().numpy().squeeze(),
+        )
 
-        self.acquire()
-        node = render_model(self.scene, model, model_out,
-                            color, "body_mesh", replace=replace)
-        self.release()
-
-        self.add_to_group("body", node)
-        return node
-
-    def render_model_with_tfs(
+    def render_model_geometry(
             self,
-            model: SMPLLayer,
-            model_out,
-            color=[1.0, 0.3, 0.3, 0.8],
             replace=True,
-            keep_pose=True,
-            render_joints=True,
-            transforms=None,
-            interpolated=False
+            pose=None,
+            name="body_mesh",
+            joints=None,
+            **kwargs
     ):
-        if model_out is None and not interpolated:
-            model_out = model()
-
-        if keep_pose:
-            node = self.get_node("body_mesh")
-            # if node is not None:
-            #original_pose = node.pose
-
-        if not interpolated:
-            self.render_joints(model_out.joints.detach().cpu().numpy().squeeze(), transforms=transforms)
-
-        self.remove_from_group("body", "body_mesh")
 
         self.acquire()
-        node = render_model_with_tfs(self.scene, model, model_out,
-                                     color, "body_mesh", replace=replace, transforms=transforms, interpolated=interpolated)
+
+        if replace:
+            self.remove_from_group("body", name)
+
+        if joints is not None:
+            self.render_joints(joints)
+
+        node = render_model_geometry(scene=self.scene, name=name, **kwargs)
+        self.add_to_group("body", node)
+
+        if pose is not None:
+            self.set_group_pose("body", pose)
         self.release()
 
-        self.add_to_group("body", node)
         return node
 
     def render_image_from_path(self, path: str, name: str = None, scale=1):

train_pose.py

@@ -153,6 +153,9 @@ def train_pose(
             loss.backward()
         return loss
 
+    # camera translation
+    R = camera.trans.detach().cpu().numpy().squeeze()
+
     # main optimization loop
     for t in range(iterations):
         loss = optimizer.step(optim_closure)
@@ -178,13 +181,13 @@ def train_pose(
             pbar.update(1)
 
         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)
-
+            renderer.render_model(
+                model=model,
+                model_out=pose_layer.cur_out,
+                transform=R
+            )
             if renderer.use_offscreen:
                 offscreen_step_output.append(renderer.get_snapshot())
-            # renderer.set_group_pose("body", R)
 
     if use_progress_bar:
         pbar.close()

utils/general.py

@@ -6,9 +6,10 @@ import cv2
 import yaml
 import os.path
 import glob
+from datetime import datetime
 
 
-def getfilename_from_conf(config, index=None):
+def getfilename_from_conf(config, index=None, include_date=True):
     """create a filename containing most training props
 
     Args:
@@ -17,6 +18,10 @@ def getfilename_from_conf(config, index=None):
     """
     name = ""
 
+    if include_date:
+        now = datetime.now()
+        name += now.strftime("%d-%m-%H-%M")
+
     if index is not None:
         name = str(index).zfill(3) + "-"
     name = name + config['pose']['optimizer']

utils/render.py

@@ -9,60 +9,33 @@ def render_model(
     scene,
     model,
     model_out,
-    color=[0.3, 0.3, 0.3, 0.8],
+    **kwargs
+):
+    return render_model_geometry(
+        scene=scene,
+        faces=model.faces,
+        vertices=model_out.vertices.detach().cpu().numpy().squeeze(),
+        **kwargs
+    )
+
+
+def render_model_geometry(
+    scene,
+    faces,
+    vertices,
+    color=[1.0, 0.3, 0.3, 0.8],
     name=None,
-    replace=False,
     pose=None
 ):
-    vertices = model_out.vertices.detach().cpu().numpy().squeeze()
-
     # set vertex colors, maybe use this to highlight accuracies
     vertex_colors = np.ones([vertices.shape[0], 4]) * color
 
     # triangulate vertex mesh
-    tri_mesh = trimesh.Trimesh(vertices, model.faces,
+    tri_mesh = trimesh.Trimesh(vertices, faces,
                                vertex_colors=vertex_colors)
 
     mesh = pyrender.Mesh.from_trimesh(tri_mesh)
 
-    if name is not None and replace:
-        for node in scene.get_nodes(name=name):
-            scene.remove_node(node)
-
-    return scene.add(mesh, name=name, pose=pose)
-
-
-def render_model_with_tfs(
-    scene,
-    model,
-    model_out,
-    color=[0.3, 0.3, 0.3, 0.8],
-    name=None,
-    replace=False,
-    pose=None,
-    transforms=None,
-    interpolated=False
-):
-
-    if not interpolated:
-        vertices = model_out.vertices.detach().cpu().numpy().squeeze()
-    else:
-        # Interpolated frames are passed as a direct array, instead of SMPLXOutput
-        vertices = model_out
-
-    # set vertex colors, maybe use this to highlight accuracies
-    vertex_colors = np.ones([vertices.shape[0], 4]) * color
-
-    # triangulate vertex mesh
-    tri_mesh = trimesh.Trimesh(vertices, model.faces,
-                               vertex_colors=vertex_colors)
-
-    mesh = pyrender.Mesh.from_trimesh(tri_mesh, poses=transforms)
-
-    if name is not None and replace:
-        for node in scene.get_nodes(name=name):
-            scene.remove_node(node)
-
     return scene.add(mesh, name=name, pose=pose)
 
 

utils/video.py

@@ -1,4 +1,5 @@
 import pickle
+from typing import Tuple
 from model import SMPLyModel
 from renderer import DefaultRenderer
 import cv2
@@ -21,14 +22,13 @@ def make_video(images, video_name: str, fps=5, ext: str = "mp4"):
     video = cv2.VideoWriter(
         video_name, fourcc, 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()
+    print("video saved to:", video_name)
 
 
 def video_from_pkl(filename, video_name, config, ext: str = "mp4"):
@@ -37,7 +37,17 @@ def video_from_pkl(filename, video_name, config, ext: str = "mp4"):
     save_to_video(model_outs, video_name, config)
 
 
-def save_to_video(poses, video_name, config, fps=30, interpolated=False):
+def save_to_video(sample_output: Tuple, video_name: str, config: object, fps=30, 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
     )
@@ -45,14 +55,32 @@ def save_to_video(poses, video_name, config, fps=30, interpolated=False):
 
     model_anim = SMPLyModel.model_from_conf(config)
 
-    frames = []
+    if interpolation_target is not None:
+        if interpolation_target % fps != 0:
+            print("[error] interpolation target must be a multiple of fps")
+            return
+        num_intermediate = int(interpolation_target / fps) - 1
+        sample_output = interpolate_poses(sample_output, num_intermediate)
 
-    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, interpolated=interpolated)
+    frames = []
+    print("[export] rendering animation frames...")
+
+    # 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_transform,
+        )
         frames.append(r.get_snapshot())
 
-    make_video(frames, video_name, fps)
+    target_fps = fps
+    if interpolation_target is not None:
+        target_fps = interpolation_target
+
+    make_video(frames, video_name, target_fps)
 
 
 def interpolate_poses(poses, num_intermediate=5):