From 4261e59f5123f6748409aa765a16d31dea127161 Mon Sep 17 00:00:00 2001
From: Roman Shapovalov <romansh@fb.com>
Date: Wed, 13 Jul 2022 05:29:09 -0700
Subject: [PATCH] Fix: making visualisation work again

Summary:
1. Respecting `visdom_show_preds` parameter when it is False.
2. Clipping the images pre-visualisation, which is important for methods like SRN that are not arare of pixel value range.

Reviewed By: bottler

Differential Revision: D37786439

fbshipit-source-id: 8dbb5104290bcc5c2829716b663cae17edc911bd
---
 projects/implicitron_trainer/experiment.py      |  2 +-
 .../visualize_reconstruction.py                 | 17 +++++++++--------
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/projects/implicitron_trainer/experiment.py b/projects/implicitron_trainer/experiment.py
index b8b6dccf..04765205 100755
--- a/projects/implicitron_trainer/experiment.py
+++ b/projects/implicitron_trainer/experiment.py
@@ -638,7 +638,7 @@ def _seed_all_random_engines(seed: int) -> None:
     random.seed(seed)
 
 
-def _setup_envvars_for_cluster(cfg) -> bool:
+def _setup_envvars_for_cluster() -> bool:
     """
     Prepares to run on cluster if relevant.
     Returns whether FAIR cluster in use.
diff --git a/projects/implicitron_trainer/visualize_reconstruction.py b/projects/implicitron_trainer/visualize_reconstruction.py
index 1e7a88a6..cfec4373 100644
--- a/projects/implicitron_trainer/visualize_reconstruction.py
+++ b/projects/implicitron_trainer/visualize_reconstruction.py
@@ -251,7 +251,7 @@ def show_predictions(
 def generate_prediction_videos(
     preds,
     sequence_name,
-    viz,
+    viz=None,
     viz_env="visualizer",
     predicted_keys=(
         "images_render",
@@ -279,19 +279,20 @@ def generate_prediction_videos(
     for rendered_pred in tqdm(preds):
         for k in predicted_keys:
             vws[k].write_frame(
-                rendered_pred[k][0].detach().cpu().numpy(),
+                rendered_pred[k][0].clip(0.0, 1.0).detach().cpu().numpy(),
                 resize=resize,
             )
 
     for k in predicted_keys:
         vws[k].get_video(quiet=True)
         print(f"Generated {vws[k].out_path}.")
-        viz.video(
-            videofile=vws[k].out_path,
-            env=viz_env,
-            win=k,  # we reuse the same window otherwise visdom dies
-            opts={"title": sequence_name + " " + k},
-        )
+        if viz is not None:
+            viz.video(
+                videofile=vws[k].out_path,
+                env=viz_env,
+                win=k,  # we reuse the same window otherwise visdom dies
+                opts={"title": sequence_name + " " + k},
+            )
 
 
 def export_scenes(