diff --git a/hydra_plugins/moai_dsl_plugin/moai_dsl_plugin.py b/hydra_plugins/moai_dsl_plugin/moai_dsl_plugin.py index 35838fb8..e7c83983 100644 --- a/hydra_plugins/moai_dsl_plugin/moai_dsl_plugin.py +++ b/hydra_plugins/moai_dsl_plugin/moai_dsl_plugin.py @@ -107,6 +107,7 @@ | "view" "(" name "," SIGNED_INT ("," SIGNED_INT)* ")" -> reshape | "transpose" "(" name "," SIGNED_INT ("," SIGNED_INT)* ")" -> transpose | "flatten" "(" name "," SIGNED_INT ["," SIGNED_INT] ")" -> flatten + | "repeat_interleave" "(" name "," SIGNED_INT "," SIGNED_INT ")" -> repeat | "zeros" "(" name ")" -> zeros_like | "ones" "(" name ")" -> ones_like | "rand" "(" name ")" -> rand_like diff --git a/moai/conf/model/monads/human/hand/mano.yaml b/moai/conf/model/monads/human/hand/mano.yaml index 59559e1a..ee5c9b69 100644 --- a/moai/conf/model/monads/human/hand/mano.yaml +++ b/moai/conf/model/monads/human/hand/mano.yaml @@ -1,4 +1,4 @@ -# @package model.monads._name_ +# @package model.monads.mano _target_: moai.monads.human.hand.MANO model_path: ??? @@ -11,4 +11,5 @@ use_betas: true use_pca: true pca_components: 12 num_betas: 10 -flat_hand_mean: false \ No newline at end of file +flat_hand_mean: false +batch_size: 1 \ No newline at end of file diff --git a/moai/conf/model/monads/math/rad2deg.yaml b/moai/conf/model/monads/math/rad2deg.yaml index fb6ce07d..623ddb6e 100644 --- a/moai/conf/model/monads/math/rad2deg.yaml +++ b/moai/conf/model/monads/math/rad2deg.yaml @@ -1,3 +1,3 @@ -# @package model.monads._name_ +# @package model.monads.rad2deg _target_: moai.monads.math.Rad2Deg \ No newline at end of file diff --git a/moai/core/execution/expression.py b/moai/core/execution/expression.py index 2ad2bb41..25b9a6b1 100644 --- a/moai/core/execution/expression.py +++ b/moai/core/execution/expression.py @@ -691,6 +691,11 @@ def flatten(self, key, *dims): # self.results.append(f'result{self.index}') # self.index += 1 + def repeat(self, key, *dims): + key = self.extract(key) + dims = list(map(int, dims)) + self._transform_operation("repeat_interleave", key, dims) + def unsqueeze(self, key, *dims): if not isinstance(key, str) or isinstance(key, Token): # NOTE: is lark.Tree key = self.extract(key) diff --git a/moai/core/model.py b/moai/core/model.py index 3e20d89f..52f27a3e 100644 --- a/moai/core/model.py +++ b/moai/core/model.py @@ -360,6 +360,13 @@ def closure(tensors, index, steps, stage, optimizer, objective): ): frequency = toolz.get(C._FREQUENCY_, iter_monitor_stage, 1) should_monitor = iter % frequency == 0 + if ( + iter_tensor_metrics := iter_monitor_stage.get(C._METRICS_, None) + ) and should_monitor: + for metric in ( + toolz.get(C._METRICS_, iter_monitor_stage, None) or [] + ): + self.named_metrics[metric](batch) if ( iter_tensor_monitor := iter_monitor_stage.get( C._MONITORS_, None @@ -369,10 +376,6 @@ def closure(tensors, index, steps, stage, optimizer, objective): toolz.get(C._FLOWS_, iter_monitor_stage, None) or [] ): self.named_flows[step](batch) - for metric in ( - toolz.get(C._METRICS_, iter_monitor_stage, None) or [] - ): - self.named_metrics[metric](batch) extras = { # TODO: step => 'lightning_step' "lightning_step": self.global_step, "epoch": self.current_epoch, @@ -462,7 +465,8 @@ def test_step( def validation_step( self, batch: typing.Dict[str, torch.Tensor], - batch_nb: int, + # batch_nb: int, + batch_idx: int, dataloader_idx: int = 0, ) -> None: batch = benedict.benedict(batch, keyattr_enabled=False) @@ -481,12 +485,18 @@ def validation_step( self.monitor, f"{C._VAL_}.{C._DATASETS_}.{dataset_name}" ) ): + extras = { + "lightning_step": self.trainer.test_loop.batch_progress.current.completed, # NOTE: self.global_step does not increment correctly + "epoch": self.current_epoch, + "batch_idx": batch_idx, + } with torch.no_grad(): for step in get_list(proc, C._FLOWS_): batch = self.named_flows[step](batch) for metric in get_list(monitor, C._METRICS_): self.named_metrics[metric](batch) - # TODO add monitors/visualization + for tensor_monitor in get_list(monitor, C._MONITORS_): + self.named_monitors[tensor_monitor](batch, extras) return batch def configure_optimizers( diff --git a/moai/data/datasets/human/amass.py b/moai/data/datasets/human/amass.py index 04c72dab..8cde16d1 100644 --- a/moai/data/datasets/human/amass.py +++ b/moai/data/datasets/human/amass.py @@ -28,11 +28,13 @@ def __init__( downsample_factor: int = 1, reconstruct: bool = False, model_type: str = "smplx", + file_type: str = None, ) -> None: super().__init__() assert_path(log, "AMASS data root path", data_root) assert_path(log, "SMPLX data path", smplx_root) - file_type = "**_poses.npz" if model_type == "smpl" else "**_stageii.npz" + if file_type is None: + file_type = "**_poses.npz" if model_type == "smpl" else "**_stageii.npz" gendered_shape_fn = "**shape.npz" if model_type == "smpl" else "**_stagei.npz" is_all_parts = isinstance(parts, str) and ("all" == parts or "**" == parts) parts = os.listdir(data_root) if is_all_parts else parts diff --git a/moai/monads/human/hand/mano.py b/moai/monads/human/hand/mano.py index c5db6bbe..f055a4d4 100644 --- a/moai/monads/human/hand/mano.py +++ b/moai/monads/human/hand/mano.py @@ -36,6 +36,7 @@ def __init__( use_pca: bool = True, flat_hand_mean: bool = False, with_tips: bool = False, + batch_size: int = 1, ): self.with_tips = with_tips super(MANO, self).__init__( @@ -47,7 +48,7 @@ def __init__( create_transl=use_translation, use_pca=use_pca, dtype=torch.float32, - batch_size=1, + batch_size=batch_size, num_pca_comps=pca_components, num_betas=num_betas, flat_hand_mean=flat_hand_mean, diff --git a/moai/serve/streaming_optimizer.py b/moai/serve/streaming_optimizer.py index e8349f4b..e5c0cad4 100644 --- a/moai/serve/streaming_optimizer.py +++ b/moai/serve/streaming_optimizer.py @@ -148,7 +148,7 @@ def handle(self, data: typing.Mapping[str, typing.Any], context: typing.Any): # get the dataloader for returned dict dataloader = td["dataloader"] # iterate over the dataloader - for batch_idx, batch in enumerate(dataloader): + for batch_idx, batch in enumerate(torch.utils.data.DataLoader(dataloader)): self.model.optimization_step = 0 # batch should be send to correct device batch = toolz.valmap(self.__to_device__, batch) diff --git a/moai/validation/metrics/generation/fid.py b/moai/validation/metrics/generation/fid.py index 52fcf55c..f2ca0449 100644 --- a/moai/validation/metrics/generation/fid.py +++ b/moai/validation/metrics/generation/fid.py @@ -14,7 +14,6 @@ def forward( self, gt: torch.Tensor, pred: torch.Tensor, - # weights: torch.Tensor=None, ) -> torch.Tensor: return {"gt": gt, "pred": pred} @@ -26,6 +25,4 @@ def compute(self, pred: np.ndarray, gt: np.ndarray) -> np.ndarray: c1, c2 = np.cov(gt[:num_samples]), np.cov(pred[:num_samples]) diff_sq = np.linalg.norm(mu1 - mu2, ord=2, axis=-1, keepdims=False) fid = diff_sq + np.trace(c1 + c2 - 2 * np.sqrt(c1 * c2)) - # if weights is not None: - # fid = fid * weights return fid diff --git a/moai/visualization/rerun/image.py b/moai/visualization/rerun/image.py index 8396ba78..eeab4c5d 100644 --- a/moai/visualization/rerun/image.py +++ b/moai/visualization/rerun/image.py @@ -31,6 +31,7 @@ def multiframe_multiview_posed_image( optimization_step: typing.Optional[int] = None, lightning_step: typing.Optional[int] = None, iteration: typing.Optional[int] = None, + jpeg_quality: int = 40, ) -> None: if optimization_step is not None: rr.set_time_sequence("optimization_step", optimization_step) @@ -58,9 +59,12 @@ def multiframe_multiview_posed_image( ), ) # log image + image = np.ascontiguousarray( + images[fr][i].transpose(-2, -1, 0) * 255 + ).astype(np.uint8) rr.log( path + f"/frame_{fr}/cam_{i}", - rr.Image(images[fr][i].transpose(-2, -1, 0)), + rr.Image(image).compress(jpeg_quality=jpeg_quality), ) @@ -72,6 +76,7 @@ def multicam_posed_image( optimization_step: typing.Optional[int] = None, lightning_step: typing.Optional[int] = None, iteration: typing.Optional[int] = None, + jpeg_quality: int = 40, ) -> None: if optimization_step is not None: rr.set_time_sequence("optimization_step", optimization_step) @@ -88,6 +93,7 @@ def multicam_posed_image( rr.Transform3D( translation=poses[i][:3, 3], mat3x3=poses[i][:3, :3], + from_parent=True, ), ) rr.log( @@ -97,9 +103,12 @@ def multicam_posed_image( ), ) # log image + image = np.ascontiguousarray(images[i].transpose(-2, -1, 0) * 255).astype( + np.uint8 + ) rr.log( path + f"/cam_{i}", - rr.Image(images[i].transpose(-2, -1, 0)), + rr.Image(image).compress(jpeg_quality=jpeg_quality), ) diff --git a/tests/dsl/test_resolver.py b/tests/dsl/test_resolver.py index 1c5c30c6..cdba60ca 100644 --- a/tests/dsl/test_resolver.py +++ b/tests/dsl/test_resolver.py @@ -357,6 +357,23 @@ def test_flatten(self, parser, highdim_tensors): x = self._parse_and_run(parser, expression, highdim_tensors) assert x.sum() == 300.0 + def test_repeat_interleave(self, parser, highdim_tensors): + expression = "repeat_interleave(fourdim, 2, 1)" + x = self._parse_and_run(parser, expression, highdim_tensors) + assert torch.equal(x, highdim_tensors["fourdim"].repeat_interleave(2, 1)) + expression = "repeat_interleave(fourdim, 2, 0)" + x = self._parse_and_run(parser, expression, highdim_tensors) + assert torch.equal(x, highdim_tensors["fourdim"].repeat_interleave(2, 0)) + expression = "repeat_interleave(single, 2, 1)" + x = self._parse_and_run(parser, expression, highdim_tensors) + assert x.sum() == 3600 * 2 # 3600 is the sum of single and 2 is the repeat + expression = "repeat_interleave(fourdim, 2, 0)" + x = self._parse_and_run(parser, expression, highdim_tensors) + assert (x[5:] != highdim_tensors["fourdim"]).sum() == 0 + expression = "repeat_interleave(fourdim, 2, 0) + ones(10, 3, 2, 6)" + x = self._parse_and_run(parser, expression, highdim_tensors) + assert x.sum() == 120 * 2 + 10 * 3 * 2 * 6 + def test_trig(self, parser, trig_tensors): expression = "sin(pi2)" x = self._parse_and_run(parser, expression, trig_tensors)