Skip to content

Commit

Permalink
Optional ground-truth depth maps in visualiser
Browse files Browse the repository at this point in the history
Summary: The code does not crash if depth map/mask are not given.

Reviewed By: bottler

Differential Revision: D45082985

fbshipit-source-id: 3610d8beb4ac897fbbe52f56a6dd012a6365b89b
  • Loading branch information
shapovalov authored and facebook-github-bot committed Apr 18, 2023
1 parent 1af6bf4 commit 0e3138e
Show file tree
Hide file tree
Showing 3 changed files with 42 additions and 36 deletions.
72 changes: 38 additions & 34 deletions pytorch3d/implicitron/evaluation/evaluate_new_view_synthesis.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -75,9 +75,11 @@ def show_depth(
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,
),
Expand All @@ -91,42 +93,44 @@ def show_depth(
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
# `TensorProperties`.
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,
Expand Down Expand Up @@ -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,
)

Expand Down
2 changes: 1 addition & 1 deletion pytorch3d/implicitron/models/generic_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -360,7 +360,7 @@ def forward(
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.
Expand Down
4 changes: 3 additions & 1 deletion tests/implicitron/test_dataset_visualize.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@

from .common_resources import get_skateboard_data

VISDOM_PORT = int(os.environ.get("VISDOM_PORT", 8097))


class TestDatasetVisualize(unittest.TestCase):
def setUp(self):
Expand Down Expand Up @@ -77,7 +79,7 @@ def setUp(self):
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
Expand Down

0 comments on commit 0e3138e

Please sign in to comment.