add TensorBoard 3D

SyneRBI · Sep 18, 2024 · 16393ed · 16393ed
1 parent 4a8a865
commit 16393ed
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -71,6 +71,7 @@ docker run --rm -it --gpus all -p 6006:6006 \
 # 3. optionally, conda/pip/apt install environment.yml/requirements.txt/apt.txt
 # 4. install metrics & run your submission
 pip install git+https://github.com/TomographicImaging/Hackathon-000-Stochastic-QualityMetrics
+pip install tensorboard-plugin-3d monai
 python petric.py &
 # 5. optionally, serve logs at <http://localhost:6006>
 tensorboard --bind_all --port 6006 --logdir ./output

diff --git a/petric.py b/petric.py
@@ -105,6 +105,12 @@ def __call__(self, algo: Algorithm):
         self.tb.add_image("coronal", np.clip(x_arr[None, :, self.coronal_slice] / self.vmax, 0, 1), algo.iteration, t)
         self.tb.add_image("sagittal", np.clip(x_arr[None, :, :, self.sagittal_slice] / self.vmax, 0, 1), algo.iteration,
                           t)
+        try:
+            from monai.visualize import plot_2d_or_3d_image
+        except ImportError:
+            pass
+        else:
+            plot_2d_or_3d_image(np.clip(x_arr[None, None] / self.vmax, 0, 1), algo.iteration, self.tb, tag='image')
         log.debug("...logged")