PR in preparation for v1.11 Release

aitviewer/aitvconfig.yaml

@@ -1,6 +1,7 @@
 # Access SMPL models.
 smplx_models: "../data/smplx_models"
 star_models: "../data/star_models"
+supr_models: "../data/supr_models"
 
 # Access to datasets.
 datasets:

aitviewer/headless.py

@@ -35,11 +35,11 @@ def __init__(self, **kwargs):
         """
         super().__init__(**kwargs)
 
-    def run(self, frame_dir=None, video_dir=None, output_fps=60):
+    def run(self, frame_dir=None, video_dir=None, output_fps=60, **export_video_kwargs):
         """Same as self.save_video, kept for backward compatibility."""
         return self.save_video(frame_dir, video_dir, output_fps)
 
-    def save_video(self, frame_dir=None, video_dir=None, output_fps=60, transparent=False):
+    def save_video(self, frame_dir=None, video_dir=None, output_fps=60, transparent=False, **export_video_kwargs):
         """
         Convenience method to run the headless rendering.
         :param frame_dir: Where to store the individual frames or None if you don't care.
@@ -51,7 +51,12 @@ def save_video(self, frame_dir=None, video_dir=None, output_fps=60, transparent=
         """
         self._init_scene()
         self.export_video(
-            output_path=video_dir, frame_dir=frame_dir, animation=True, output_fps=output_fps, transparent=transparent
+            output_path=video_dir,
+            frame_dir=frame_dir,
+            animation=True,
+            output_fps=output_fps,
+            transparent=transparent,
+            **export_video_kwargs,
         )
 
     def save_frame(self, file_path, scale_factor: float = None):

aitviewer/models/supr.py

@@ -0,0 +1,143 @@
+"""
+Copyright (C) 2023  ETH Zurich, Manuel Kaufmann, Velko Vechev, Dario Mylonopoulos
+
+This program is free software: you can redistribute it and/or modify
+it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+This program is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with this program.  If not, see <https://www.gnu.org/licenses/>.
+"""
+import collections
+import os
+
+import torch
+import torch.nn.functional as F
+
+try:
+    from supr.config import cfg
+    from supr.pytorch.supr import SUPR
+except Exception as e:
+    raise ImportError(f"Cannot import SUPR. Please run `pip install git+https://github.com/ahmedosman/SUPR.git`\n{e}")
+
+from aitviewer.configuration import CONFIG as C
+from aitviewer.utils.so3 import aa2rot_torch as aa2rot
+from aitviewer.utils.so3 import rot2aa_torch as rot2aa
+
+
+class SUPRLayer(SUPR):
+    """Wraps the publicly available SUPR model to match SMPLX model interface"""
+
+    def __init__(self, gender="male", num_betas=10, constrained=False, device=None, dtype=None):
+        """
+        Initializer.
+        :param gender: Which gender to load.
+        :param num_betas: Number of shape components.
+        :param device: CPU or GPU.
+        :param dtype: The pytorch floating point data type.
+        """
+        # Configure SUPR model before initializing
+        cfg.data_type = dtype if dtype is not None else C.f_precision
+        path_model = os.path.join(C.supr_models, f'supr_{gender}{"_constrained" if constrained else ""}.npy')
+        super(SUPRLayer, self).__init__(path_model, num_betas=num_betas)
+
+        self.device = device if device is not None else C.device
+        self.model_type = "supr"
+        self._parents = None
+        self._children = None
+
+    @property
+    def parents(self):
+        """Return how the joints are connected in the kinematic chain where parents[i, 0] is the parent of
+        joint parents[i, 1]."""
+        if self._parents is None:
+            self._parents = self.kintree_table.transpose(0, 1).cpu().numpy()
+        return self._parents
+
+    @property
+    def joint_children(self):
+        """Return the children of each joint in the kinematic chain."""
+        if self._children is None:
+            self._children = collections.defaultdict(list)
+            for bone in self.parents:
+                if bone[0] != -1:
+                    self._children[bone[0]].append(bone[1])
+        return self._children
+
+    def skeletons(self):
+        """Return how the joints are connected in the kinematic chain where skeleton[0, i] is the parent of
+        joint skeleton[1, i]."""
+        kintree_table = self.kintree_table
+        kintree_table[:, 0] = -1
+        return {
+            "all": kintree_table,
+            "body": kintree_table[:, : self.n_joints_body + 1],
+        }
+
+    @property
+    def n_joints_body(self):
+        return self.parent.shape[0]
+
+    @property
+    def n_joints_total(self):
+        return self.n_joints_body + 1
+
+    def forward(self, poses_body, betas=None, poses_root=None, trans=None, normalize_root=False):
+        """
+        forwards the model
+        :param poses_body: Pose parameters.
+        :param poses_root: Pose parameters for the root joint.
+        :param beta: Beta parameters.
+        :param trans: Root translation.
+        :param normalize_root: Makes poses relative to the root joint (useful for globally rotated captures).
+        :return: Deformed surface vertices, transformed joints
+        """
+        poses, betas, trans = self.preprocess(poses_body, betas, poses_root, trans, normalize_root)
+
+        # SUPR repo currently hardcodes floats.
+        v = super().forward(pose=poses.float(), betas=betas.float(), trans=trans.float())
+        J = v.J_transformed
+        return v, J
+
+    def preprocess(self, poses_body, betas=None, poses_root=None, trans=None, normalize_root=False):
+        batch_size = poses_body.shape[0]
+
+        if poses_root is None:
+            poses_root = torch.zeros([batch_size, 3]).to(dtype=poses_body.dtype, device=self.device)
+        if trans is None:
+            # If we don't supply the root translation explicitly, it falls back to using self.bm.trans
+            # which might not be zero since it is a trainable param that can get updated.
+            trans = torch.zeros([batch_size, 3]).to(dtype=poses_body.dtype, device=self.device)
+
+        if normalize_root:
+            # Make everything relative to the first root orientation.
+            root_ori = aa2rot(poses_root)
+            first_root_ori = torch.inverse(root_ori[0:1])
+            root_ori = torch.matmul(first_root_ori, root_ori)
+            poses_root = rot2aa(root_ori)
+            trans = torch.matmul(first_root_ori.unsqueeze(0), trans.unsqueeze(-1)).squeeze()
+            trans = trans - trans[0:1]
+
+        poses = torch.cat((poses_root, poses_body), dim=1)
+
+        if betas is None:
+            betas = torch.zeros([batch_size, self.num_betas]).to(dtype=poses_body.dtype, device=self.device)
+
+        # Batch shapes if they don't match batch dimension.
+        if betas.shape[0] != batch_size:
+            betas = betas.repeat(batch_size, 1)
+
+        # Lower bound betas
+        if betas.shape[1] < self.num_betas:
+            betas = torch.nn.functional.pad(betas, [0, self.num_betas - betas.shape[1]])
+
+        # Upper bound betas
+        betas = betas[:, : self.num_betas]
+
+        return poses, betas, trans

aitviewer/renderables/billboard.py

@@ -14,13 +14,15 @@
 You should have received a copy of the GNU General Public License
 along with this program.  If not, see <https://www.gnu.org/licenses/>.
 """
+import os
 import pickle
 from typing import List, Union
 
 import cv2
 import moderngl
 import numpy as np
 from moderngl_window.opengl.vao import VAO
+from pxr import Gf, Sdf, UsdGeom
 from trimesh.triangles import points_to_barycentric
 
 from aitviewer.scene.camera import Camera, OpenCVCamera
@@ -30,6 +32,7 @@
     get_outline_program,
     get_screen_texture_program,
 )
+from aitviewer.utils import usd
 from aitviewer.utils.decorators import hooked
 
 
@@ -298,3 +301,59 @@ def gui_material(self, imgui, show_advanced=True):
             1.0,
             "%.2f",
         )
+
+    def export_usd(self, stage, usd_path: str, directory: str = None, verbose=False):
+        name = f"{self.name}_{self.uid:03}".replace(" ", "_")
+        usd_path = f"{usd_path}/{name}"
+
+        if not directory:
+            if isinstance(self.textures, list) and isinstance(self.textures[0], str):
+                if self.textures[0].endswith((".pickle", ".pkl")):
+                    print(
+                        f"Failed to export billboard: {self.name}. Textures must not be pickle files when not exporting to a directory."
+                    )
+                    return
+            else:
+                print(
+                    f"Failed to export billboard: {self.name}. Textures must be paths to image files when not exporting to a directory."
+                )
+                return
+
+        # Transform.
+        xform = UsdGeom.Xform.Define(stage, usd_path)
+        a_xform = xform.AddTransformOp()
+        a_xform.Set(Gf.Matrix4d(self.get_local_transform().astype(np.float64).T))
+
+        # Geometry.
+        mesh = UsdGeom.Mesh.Define(stage, usd_path + "/" + self.name.replace(" ", "_"))
+        a_vertices = mesh.CreatePointsAttr()
+        for i in range(self.vertices.shape[0]):
+            a_vertices.Set(time=i + 1, value=self.vertices[i])
+        mesh.CreateFaceVertexCountsAttr(np.array([4]))
+        mesh.CreateFaceVertexIndicesAttr([0, 1, 3, 2])
+        a_uv = UsdGeom.PrimvarsAPI(mesh).CreatePrimvar(
+            "st", Sdf.ValueTypeNames.TexCoord2fArray, UsdGeom.Tokens.faceVarying
+        )
+        a_uv.Set(time=1, value=np.array([[1.0, 1.0], [1.0, 0.0], [0.0, 0.0], [0.0, 1.0]]))
+
+        # Textures.
+        if not directory:
+            usd.add_texture(stage, mesh, usd_path, os.path.abspath(self.textures[0]))
+        else:
+            if isinstance(self.textures, list) and isinstance(self.textures[0], str):
+                for i, path in enumerate(self.textures):
+                    if path.endswith((".pickle", "pkl")):
+                        img = pickle.load(open(path, "rb"))
+                        texture_path = usd.save_image_as_texture(img, f"img_{i:03}.png", name, directory)
+                    else:
+                        texture_path = usd.copy_texture(path, name, directory)
+                    if i == 0:
+                        usd.add_texture(stage, mesh, usd_path, texture_path)
+            else:
+                for i in range(len(self.textures)):
+                    img = self.textures[i]
+                    texture_path = usd.save_image_as_texture(img, f"img_{i:03}.png", name, directory)
+                    if i == 0:
+                        usd.add_texture(stage, mesh, usd_path, texture_path)
+
+        self._export_usd_recursively(stage, usd_path, directory, verbose)

aitviewer/renderables/lines.py

@@ -19,6 +19,7 @@
 import trimesh
 from moderngl_window.opengl.vao import VAO
 
+from aitviewer.renderables.spheres import SpheresTrail
 from aitviewer.scene.material import Material
 from aitviewer.scene.node import Node
 from aitviewer.shaders import (
@@ -588,3 +589,25 @@ def render(self, camera, **kwargs):
     def release(self):
         if self.is_renderable:
             self.vao.release()
+
+
+class LinesTrail(Lines):
+    """A sequence of lines that leave a trail, i.e. the past lines keep being rendered."""
+
+    def __init__(self, lines, with_spheres=True, r_base=0.01, r_tip=None, color=(0.0, 0.0, 1.0, 1.0), **kwargs):
+        if "mode" not in kwargs:
+            kwargs["mode"] = "line_strip"
+        super().__init__(lines, r_base, r_tip, color, **kwargs)
+        self.n_frames = lines.shape[0]
+
+        if with_spheres:
+            spheres_trail = SpheresTrail(lines, radius=r_base * 4.0, color=color)
+            self._add_node(spheres_trail, enabled=True, show_in_hierarchy=True)
+
+    def on_frame_update(self):
+        self.n_lines = self.current_frame_id
+        super().on_frame_update()
+
+    def make_renderable(self, ctx):
+        super().make_renderable(ctx)
+        self.on_frame_update()