Merge branch 'main' of https://github.com/moverseai/moai

docker/serve/azure/pt220-cu118/Dockerfile

@@ -1,4 +1,4 @@
-FROM ghcr.io/moverseai/moai/serve:v1.0.0-pt220-cu118
+FROM ghcr.io/moverseai/moai/serve:v1.1.0-pt220-cu118
 
 USER root
 

moai/conf/data/train/iterator/window.yaml

@@ -0,0 +1,6 @@
+# @package data.train.iterator
+
+_target_: moai.data.iterator.Windowed
+window_size: ???
+stride: 1
+internal_stride: 1

moai/conf/engine/modules/rerun.yaml

@@ -11,3 +11,6 @@ world_coordinates: RUF
 add_floor: true
 root: /
 memory_limit: 75% #NOTE: can also be absolute, `4GB` for example
+log:
+  handler: moai/logs # null
+  level: ${hydra:hydra_logging.root.level}

moai/conf/examples/smplifyx/flows.yaml

@@ -9,16 +9,28 @@ _moai_:
         _out_: [init_trans]
       preprocess:
         embedding:
-          void: [color]
+          void:
+            - ${mi:"color[0:1, 1:2, :, :]"}
           _out_: [embedding]
         translation:
-          void: [color]
+          void:
+            - ${mi:"sq(color, 0)"}
           _out_: [translation]
         betas:
-          void: [color]
+          void:
+            - ${mi:"flatten(color, 1)"}
           _out_: [betas]
+        _mi_1:
+          expression:
+            - ${mi:"embedding + zeros(embedding)"}
+            - ${mi:"embedding_1 / ones(embedding)"}
+            - ${mi:"flatten(embedding_2, 1)"}
+            - ${mi:"sq(embedding_3, 0)"}
+            - ${mi:"view(embedding_4, 1, 32)"}
+          _out_:
+            [embedding_1, embedding_2, embedding_3, embedding_4, embedding_5]
         vposer1: # vposer2:
-          decode: [embedding]
+          decode: [embedding_5] # [embedding]
           _out_: [decoded]
         smplx:
           pose: [decoded.pose]

moai/conf/model/monads/geometry/projection/mv_weak_perspective_camera.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.mv_weak_perspective_camera
 
 _target_: moai.monads.geometry.MVWeakPerspective
 # focal_length: 5000.0

moai/conf/model/monads/geometry/projection/project.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.project
 
 _target_: moai.monads.geometry.Projection
 xyz_at: channel # one of ['channel', 'row]

moai/conf/model/monads/geometry/rotation/conversion.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.conversion
 
 _target_: moai.monads.geometry.rotations.ConvertRotation
 src: ??? # from representation

moai/conf/model/monads/geometry/transformation/transform.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.transform
 
 _target_: moai.monads.geometry.Transformation
 xyz_in_at: channel  # one of ['channel', 'row]

moai/conf/model/monads/networks/autoencoder.yaml

@@ -1,4 +1,4 @@
 # @package model.monads.autoencoder
 
-_target_: moai.monads.networks.autoencoder.AutoEncoder
+_target_: moai.monads.networks._factory.autoencoder.AutoEncoder
 checkpoint: ???

moai/conf/model/monads/tensor/stack.yaml

@@ -1,4 +1,4 @@
-# @package model.monads._name_
+# @package model.monads.stack
 
 _target_: moai.monads.tensor.Stack
 dim: 0

moai/conf/model/parameters/groups/human/body/retargetting.yaml

@@ -0,0 +1,7 @@
+# @package model.parameters.groups.retargetting
+
+_target_: moai.parameters.selectors.model.ModelParameterSelector
+monads:
+  - main_flow.pose_params
+  - main_flow.translation_params
+force_grad: true

moai/engine/modules/rerun.py

@@ -1,10 +1,13 @@
+import logging
 import typing
 
 import colour
 import numpy as np
 import toolz
+from omegaconf.omegaconf import DictConfig
 
 from moai.utils.color.colormap import random_color
+from moai.utils.funcs import get
 
 try:
     import rerun as rr
@@ -16,6 +19,15 @@
 __all__ = ["Rerun"]
 
 
+class RerunFilter(logging.Filter):
+    def __init__(self) -> None:
+        super().__init__("rerun")
+
+    def filter(self, record):
+        record.msg = record.msg.removeprefix(":moai: ")
+        return record
+
+
 class Rerun:
     r"""
     This logger should support all the different types of rr logging.
@@ -39,6 +51,7 @@ def __init__(
         add_floor: bool = True,
         root: str = "/",
         memory_limit: str = "75%",
+        log: DictConfig = {},
     ) -> None:
         # NOTE: https://github.com/iterative/dvc/issues/9731
         rr.init(name)
@@ -62,6 +75,15 @@ def __init__(
         if add_floor:
             self._create_floor(root)
         # NOTE: add more global gizmos (e.g. tr-axis)
+        if handler_name := get(log, "handler"):
+            handler = rr.LoggingHandler(handler_name)
+            handler.addFilter(RerunFilter())
+            logging.getLogger().addHandler(handler)
+            if level := get(log, "level"):
+                if isinstance(level, str):
+                    level = level.upper()
+                level = logging.getLevelName(level)
+                logging.getLogger().setLevel(level)
 
     def _create_scalar_plots(self, root: str, plots) -> None:
         for plot in plots:

moai/engine/progressbar.py

@@ -1,4 +1,5 @@
 import rich.progress
+from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks import RichProgressBar
 from pytorch_lightning.callbacks.progress.rich_progress import RichProgressBarTheme
 
@@ -22,6 +23,8 @@
 # NOTE: check https://github.com/Textualize/rich/discussions/482
 # NOTE: check https://github.com/facebookresearch/EGG/blob/a139946a73d45553360a7f897626d1ae20759f12/egg/core/callbacks.py#L335
 # NOTE: check https://github.com/Textualize/rich/discussions/921
+
+
 class MoaiProgressBar(RichProgressBar):
     def __init__(self) -> None:
         super().__init__(
@@ -39,7 +42,7 @@ def __init__(self) -> None:
     #         CustomTimeColumn(style=self.theme.time),
     #         ProcessingSpeedColumn(style=self.theme.processing_speed),
     #     ]
-    def configure_columns(self, trainer: "pl.Trainer") -> list:
+    def configure_columns(self, trainer: Trainer) -> list:
         original = super().configure_columns(trainer)
         moai_column = rich.progress.TextColumn(":moai:")
         spinner_column = rich.progress.SpinnerColumn(spinner_name="dots5")

moai/monads/generation/tensor/torch.py

@@ -2,6 +2,7 @@
 
 import numpy as np
 import omegaconf.omegaconf
+import toolz
 import torch
 
 __all__ = [
@@ -15,6 +16,7 @@
     "ZerosLike",
     "RandomLike",
     "OnesLike",
+    "TemporalParams",
 ]
 
 
@@ -166,12 +168,17 @@ def __init__(
         init: str = "zeros",  # one of [zeros, ones, rand, randn],
     ):
         super(Parameter, self).__init__()
-        self.register_parameter(
-            "value",
-            torch.nn.Parameter(
-                getattr(torch, init)(tuple(shape))
-            ),  # TODO: check omegaconf's convert type annotation
-        )
+        if init == "eye":
+            self.register_parameter(
+                "value", torch.nn.Parameter(torch.eye(3).repeat(*shape))
+            )
+        else:
+            self.register_parameter(
+                "value",
+                torch.nn.Parameter(
+                    getattr(torch, init)(tuple(shape))
+                ),  # TODO: check omegaconf's convert type annotation
+            )
 
     def forward(self, void: torch.Tensor) -> torch.nn.parameter.Parameter:
         return self.value
@@ -193,3 +200,37 @@ def __init__(
 
     def forward(self, void: torch.Tensor) -> torch.nn.parameter.Parameter:
         return dict(self.named_parameters())
+
+
+class TemporalParams(torch.nn.Module):
+    r"""
+    Base class for temporal parameter generation.
+
+    Args:
+        parameters (omegaconf.DictConfig): parameters for temporal params
+        window_size (int): number of frames to generate params for
+    """
+
+    def __init__(
+        self,
+        parameters: omegaconf.DictConfig,
+        window_size: int,
+    ) -> None:
+        super().__init__()
+        window_size = int(window_size) if isinstance(window_size, str) else window_size
+        for name, param in parameters.items():
+            # each param should be generated for each frame in the window
+            for i in range(window_size):
+                self.register_parameter(
+                    str(name) + str(i),
+                    torch.nn.Parameter(
+                        getattr(torch, toolz.get("init", param, "zeros"))(
+                            tuple(param.shape)
+                        )
+                    ),  # TODO: check omegaconf's convert type annotation
+                )
+
+        self.window_size = window_size
+
+    def forward(self, void: torch.Tensor) -> torch.nn.parameter.Parameter:
+        return dict(self.named_parameters())

moai/monads/human/pose/forward_kinematics.py

@@ -1,3 +1,4 @@
+import logging
 import typing
 
 import numpy as np
@@ -6,25 +7,35 @@
 __all__ = ["ForwardKinematics"]
 
 
+log = logging.getLogger(__name__)
+
+
 class ForwardKinematics(torch.nn.Module):
     def __init__(
-        self, parents
+        self,
+        parents: list,
+        offsets: list = None,
     ):  # TODO: add a col/row major param to adjust offset slicing
         super().__init__()
-        self.parents = parents  # TODO: register buffer from list?
+        self.register_buffer("parents", torch.Tensor(parents).int())
+        if offsets is not None:
+            log.warning("offset not set, should be passed in the forward method")
+            self.register_buffer("offsets", torch.Tensor(offsets))
 
     def forward(
         self,  # TODO: add parents tensor input?
         rotation: torch.Tensor,  # [B, (T), J, 3, 3]
         position: torch.Tensor,  # [B, (T), 3]
-        offset: torch.Tensor,  # [B, (T), J, 3]
+        offset: typing.Optional[torch.Tensor] = None,  # [B, (T), J, 3]
         parents: typing.Optional[torch.Tensor] = None,  # [B, J]
     ) -> typing.Dict[str, torch.Tensor]:  # { [B, (T), J, 3], [B, (T), J, 3, 3] }
         joints = torch.empty(rotation.shape[:-1], device=rotation.device)
         joints[..., 0, :] = position.clone()  # first joint according to global position
-        offset = offset[
-            :, np.newaxis, ..., np.newaxis
-        ]  # NOTE: careful, col vs row major order
+        offset = (
+            offset[:, np.newaxis, ..., np.newaxis]
+            if offset is not None
+            else self.offsets[:, np.newaxis, ..., np.newaxis]
+        )  # NOTE: careful, col vs row major order
         # offset = offset[np.newaxis, :, np.newaxis, :] #NOTE: careful, col vs row major order
         global_rotation = rotation.clone()
         # global_rotation = torch.empty(rotation.shape, device=rotation.device)
@@ -33,7 +44,7 @@ def forward(
         parent_indices = (
             parents[0].detach().cpu()
             if parents is not None
-            else self.parents[0].detach().cpu()
+            else (self.parents[0].detach().cpu())
         )
         if (
             parent_indices.shape[-1] == offset.shape[-3]

moai/serve/handlers/azure/blob.py

@@ -47,10 +47,9 @@ def __call__(
     ) -> typing.Any:
 
         if self.json_key not in json:
-            log.error(f"json key: {self.json_key}, not found in json request")        
+            log.error(f"json key: {self.json_key}, not found in json request")
         working_dir = json[self.json_key]
 
-
         # initialize connection to Azure Blob Storage
         connect_str = json[self.connection_string]
         try:
@@ -103,7 +102,7 @@ def __init__(
         # self.blob_service_client = BlobServiceClient.from_connection_string(
         #     connection_string,
         # )
- 
+
         self.connection_string = connection_string
         self.container_name = container_name
         self.blob_paths = blob_paths
@@ -122,9 +121,9 @@ def __call__(
         # NOTE: void is the input json response
         # TODO: need to check batched inference
         input_json = void[0].get("body") or void[0].get("raw")
-        
+
         if self.json_key not in input_json:
-            log.error(f"json key: {self.json_key}, not found in json request")        
+            log.error(f"json key: {self.json_key}, not found in json request")
         working_dir = input_json[self.json_key]
 
         # initialize connection to Azure Blob Storage

moai/validation/metrics/common/distance_accuracy.py

@@ -1,10 +1,12 @@
 import numpy as np
 import torch
 
+from moai.validation.metric import MoaiMetric
+
 __all__ = ["DistanceAccuracy"]
 
 
-class DistanceAccuracy(torch.nn.Module):
+class DistanceAccuracy(MoaiMetric):
     def __init__(self, threshold: float = 0.05) -> None:
         super().__init__()
         self.threshold = threshold
@@ -15,3 +17,6 @@ def forward(
         total = np.prod(list(distance.shape)[1:])
         acc = (distance.flatten(1) < self.threshold).sum(1) / total
         return acc.mean()
+
+    def compute(self, accs: np.ndarray) -> np.ndarray:
+        return accs.mean()