From 0e3138eca89336d222d9338687741da824a39443 Mon Sep 17 00:00:00 2001 From: Roman Shapovalov Date: Tue, 18 Apr 2023 07:00:17 -0700 Subject: [PATCH] Optional ground-truth depth maps in visualiser Summary: The code does not crash if depth map/mask are not given. Reviewed By: bottler Differential Revision: D45082985 fbshipit-source-id: 3610d8beb4ac897fbbe52f56a6dd012a6365b89b --- .../evaluation/evaluate_new_view_synthesis.py | 72 ++++++++++--------- pytorch3d/implicitron/models/generic_model.py | 2 +- tests/implicitron/test_dataset_visualize.py | 4 +- 3 files changed, 42 insertions(+), 36 deletions(-) diff --git a/pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py b/pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py index e380208b..1ee1a779 100644 --- a/pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py +++ b/pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py @@ -38,8 +38,8 @@ class _Visualizer: image_render: torch.Tensor image_rgb_masked: torch.Tensor depth_render: torch.Tensor - depth_map: torch.Tensor - depth_mask: torch.Tensor + depth_map: Optional[torch.Tensor] + depth_mask: Optional[torch.Tensor] visdom_env: str = "eval_debug" @@ -75,9 +75,11 @@ class _Visualizer: viz = self._viz viz.images( torch.cat( - ( - make_depth_image(self.depth_render, loss_mask_now), - make_depth_image(self.depth_map, loss_mask_now), + (make_depth_image(self.depth_render, loss_mask_now),) + + ( + (make_depth_image(self.depth_map, loss_mask_now),) + if self.depth_map is not None + else () ), dim=3, ), @@ -91,12 +93,13 @@ class _Visualizer: win="depth_abs" + name_postfix + "_mask", opts={"title": f"depth_abs_{name_postfix}_{depth_loss:1.2f}_mask"}, ) - viz.images( - self.depth_mask, - env=self.visdom_env, - win="depth_abs" + name_postfix + "_maskd", - opts={"title": f"depth_abs_{name_postfix}_{depth_loss:1.2f}_maskd"}, - ) + if self.depth_mask is not None: + viz.images( + self.depth_mask, + env=self.visdom_env, + win="depth_abs" + name_postfix + "_maskd", + opts={"title": f"depth_abs_{name_postfix}_{depth_loss:1.2f}_maskd"}, + ) # show the 3D plot # pyre-fixme[9]: viewpoint_trivial has type `PerspectiveCameras`; used as @@ -104,29 +107,30 @@ class _Visualizer: viewpoint_trivial: PerspectiveCameras = PerspectiveCameras().to( loss_mask_now.device ) - pcl_pred = get_rgbd_point_cloud( - viewpoint_trivial, - self.image_render, - self.depth_render, - # mask_crop, - torch.ones_like(self.depth_render), - # loss_mask_now, - ) - pcl_gt = get_rgbd_point_cloud( - viewpoint_trivial, - self.image_rgb_masked, - self.depth_map, - # mask_crop, - torch.ones_like(self.depth_map), - # loss_mask_now, - ) _pcls = { - pn: p - for pn, p in zip(("pred_depth", "gt_depth"), (pcl_pred, pcl_gt)) - if int(p.num_points_per_cloud()) > 0 + "pred_depth": get_rgbd_point_cloud( + viewpoint_trivial, + self.image_render, + self.depth_render, + # mask_crop, + torch.ones_like(self.depth_render), + # loss_mask_now, + ) } + if self.depth_map is not None: + _pcls["gt_depth"] = get_rgbd_point_cloud( + viewpoint_trivial, + self.image_rgb_masked, + self.depth_map, + # mask_crop, + torch.ones_like(self.depth_map), + # loss_mask_now, + ) + + _pcls = {pn: p for pn, p in _pcls.items() if int(p.num_points_per_cloud()) > 0} + plotlyplot = plot_scene( - {f"pcl{name_postfix}": _pcls}, + {f"pcl{name_postfix}": _pcls}, # pyre-ignore camera_scale=1.0, pointcloud_max_points=10000, pointcloud_marker_size=1, @@ -277,10 +281,10 @@ def eval_batch( image_render=image_render, image_rgb_masked=image_rgb_masked, depth_render=cloned_render["depth_render"], - # pyre-fixme[6]: Expected `Tensor` for 4th param but got - # `Optional[torch.Tensor]`. depth_map=frame_data.depth_map, - depth_mask=frame_data.depth_mask[:1], + depth_mask=frame_data.depth_mask[:1] + if frame_data.depth_mask is not None + else None, visdom_env=visualize_visdom_env, ) diff --git a/pytorch3d/implicitron/models/generic_model.py b/pytorch3d/implicitron/models/generic_model.py index 6e336105..869fd488 100644 --- a/pytorch3d/implicitron/models/generic_model.py +++ b/pytorch3d/implicitron/models/generic_model.py @@ -360,7 +360,7 @@ class GenericModel(ImplicitronModelBase): # pyre-ignore: 13 and source images, which will be used for intersecting with target rays. fg_probability: A tensor of shape `(B, 1, H, W)` containing a batch of foreground masks. - mask_crop: A binary tensor of shape `(B, 1, H, W)` deonting valid + mask_crop: A binary tensor of shape `(B, 1, H, W)` denoting valid regions in the input images (i.e. regions that do not correspond to, e.g., zero-padding). When the `RaySampler`'s sampling mode is set to "mask_sample", rays will be sampled in the non zero regions. diff --git a/tests/implicitron/test_dataset_visualize.py b/tests/implicitron/test_dataset_visualize.py index c7b66d8b..d7af5b03 100644 --- a/tests/implicitron/test_dataset_visualize.py +++ b/tests/implicitron/test_dataset_visualize.py @@ -26,6 +26,8 @@ from tests.common_testing import interactive_testing_requested from .common_resources import get_skateboard_data +VISDOM_PORT = int(os.environ.get("VISDOM_PORT", 8097)) + class TestDatasetVisualize(unittest.TestCase): def setUp(self): @@ -77,7 +79,7 @@ class TestDatasetVisualize(unittest.TestCase): for k, dataset in self.datasets.items() } ) - self.visdom = Visdom() + self.visdom = Visdom(port=VISDOM_PORT) if not self.visdom.check_connection(): print("Visdom server not running! Disabling visdom visualizations.") self.visdom = None