add video renderer

example_fit.py

@@ -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)
+make_video(step_imgs, "test.avi")
-ax.plot(train_loss[1::], label='sgd')
+
-ax.set(xlabel="Training iteration", ylabel="Loss", title='Training loss')
+# color = r.get_snapshot()
-fig.savefig("results/" + name + ".png")
+# plt.imshow(color)
-ax.legend()
+# plt.show()
-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()

example_fit_anim.py

@@ -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:
+    for idx in dataset:
-        keypoints, confidence, img_path = get_next_frame(idx)
 
-        est_scale = estimate_scale(joints, keypoints)
+        init_keypoints, init_joints, keypoints, conf, est_scale, r, img_path = setup_training(
-
+            model=model,
-        # apply scaling to keypoints
+            renderer=True,
-        keypoints = keypoints * est_scale
+            offscreen=True,
-
+            dataset=dataset,
-        init_joints = get_torso(joints)
+            sample_index=idx
-        init_keypoints = get_torso(keypoints)
+        )
 
         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.
 '''

renderer.py

@@ -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,
+        if self.use_offscreen:
-            run_in_thread=run_in_thread,
+            self.offscreen = pyrender.OffscreenRenderer(
-            use_reymond_lighting=use_reymond_lighting,
+                # self.scene,
-            viewport_size=tuple(d // 2 for d in viewport_size)
+                # 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):
+    def render_image(self, image, scale, name="source_img"):
-        _ = render_image_plane(self.scene, image, scale, name)
+        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

train_pose.py

@@ -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
     )

utils/general.py

@@ -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,

utils/render.py

@@ -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(