fix optional savefig code

agimus-project · Feb 15, 2024 · 9eb0e4c · 9eb0e4c
1 parent 719de9b
commit 9eb0e4c
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 20 deletions.
diff --git a/tests/test_renderer_bullet.py b/tests/test_renderer_bullet.py
@@ -1,4 +1,4 @@
-"""Set of unit tests for Panda3D renderer."""
+"""Set of unit tests for bullet renderer."""
 import unittest
 from pathlib import Path
 
@@ -14,7 +14,7 @@
 
 
 class TestBulletRenderer(unittest.TestCase):
-    """Unit tests for Panda3D renderer."""
+    """Unit tests for bullet renderer."""
 
     def setUp(self) -> None:
 
@@ -84,11 +84,13 @@ def test_scene_renderer(self):
 
         if SAVEFIG:
             import matplotlib.pyplot as plt
-            plt.subplot(1,2,1)
+            plt.subplot(2,2,1)
             plt.imshow(rgb)
-            plt.subplot(1,2,2)
+            plt.subplot(2,2,3)
             plt.imshow(depth, cmap=plt.cm.gray_r)
-            fig_path = self.obj_path.parent / f'panda3d_{self.obj_label}_render.png'
+            plt.subplot(2,2,4)
+            plt.imshow(binary_mask, cmap=plt.cm.gray)
+            fig_path = self.obj_path.parent / f'bullet_{self.obj_label}_scene_render.png'
             print(f'Saved {fig_path}')
             plt.savefig(fig_path)
 
@@ -177,11 +179,13 @@ def test_batch_renderer(self):
 
         if SAVEFIG:
             import matplotlib.pyplot as plt
-            plt.subplot(1,2,1)
+            plt.subplot(2,2,1)
             plt.imshow(rgb)
-            plt.subplot(1,2,2)
+            plt.subplot(2,2,3)
             plt.imshow(depth, cmap=plt.cm.gray_r)
-            fig_path = self.obj_path.parent / f'panda3d_{self.obj_label}_render.png'
+            plt.subplot(2,2,4)
+            plt.imshow(binary_mask, cmap=plt.cm.gray)
+            fig_path = self.obj_path.parent / f'bullet_{self.obj_label}_batch_render.png'
             print(f'Saved {fig_path}')
             plt.savefig(fig_path)
 

diff --git a/tests/test_renderer_panda3d.py b/tests/test_renderer_panda3d.py
@@ -121,13 +121,15 @@ def test_scene_renderer(self):
 
         if SAVEFIG:
             import matplotlib.pyplot as plt
-            plt.subplot(1,3,1)
+            plt.subplot(2,2,1)
             plt.imshow(rgb)
-            plt.subplot(1,3,2)
-            plt.imshow(depth, cmap=plt.cm.gray_r)
-            plt.subplot(1,3,3)
+            plt.subplot(2,2,2)
             plt.imshow(normals)
-            fig_path = self.obj_path.parent / f'panda3d_{self.obj_label}_render.png'
+            plt.subplot(2,2,3)
+            plt.imshow(depth, cmap=plt.cm.gray_r)
+            plt.subplot(2,2,4)
+            plt.imshow(binary_mask, cmap=plt.cm.gray)
+            fig_path = self.obj_path.parent / f'panda3d_{self.obj_label}_scene_render.png'
             print(f'Saved {fig_path}')
             plt.savefig(fig_path)
 
@@ -262,17 +264,17 @@ def test_batch_renderer(self):
         depth = renderings.depths[0].movedim(0,-1).numpy()  # (Nc,3,h,w) -> (h,w,3)
         binary_mask = renderings.binary_masks[0].movedim(0,-1).numpy()  # (Nc,1,h,w) -> (h,w,1)
 
-
-
         if SAVEFIG:
             import matplotlib.pyplot as plt
-            plt.subplot(1,3,1)
+            plt.subplot(2,2,1)
             plt.imshow(rgb)
-            plt.subplot(1,3,2)
-            plt.imshow(depth, cmap=plt.cm.gray_r)
-            plt.subplot(1,3,3)
+            plt.subplot(2,2,2)
             plt.imshow(normals)
-            fig_path = self.obj_path.parent / f'panda3d_{self.obj_label}_render.png'
+            plt.subplot(2,2,3)
+            plt.imshow(depth, cmap=plt.cm.gray_r)
+            plt.subplot(2,2,4)
+            plt.imshow(binary_mask, cmap=plt.cm.gray)
+            fig_path = self.obj_path.parent / f'panda3d_{self.obj_label}_batch_render.png'
             print(f'Saved {fig_path}')
             plt.savefig(fig_path)