changed tests to save rendered images as pngs

Author: albertaillet

.gitignore

@@ -3,3 +3,4 @@ colab_*.py
 *.pyc
 assets/other/
 .vscode/
+tmp/

tests.py

@@ -1,49 +1,45 @@
 from sys import argv
+from time import time
+from pathlib import Path
+from numpy import hstack
 from yaml import safe_load
-from utils.plot import to_rgb
 from typeguard import check_type
-from matplotlib import pyplot as plt
+from utils.plot import to_rgb, fromarray, Image
 from raymarch import render_scene, RenderedImages
 from builder import build_scene, check_scene_dict, SceneDict
 from app import SCENES_PATH
 
 
-def test_all_scenes():
+def test_build_scene(scene_path: Path) -> None:
     annotations = render_scene.__annotations__
     checked_args = ('objects', 'camera', 'view_size', 'light_dir')
-    for path in SCENES_PATH.glob('*.yml'):
-        scene_dict = check_type(check_scene_dict(safe_load(open(path))), SceneDict)
-        scene = build_scene(scene_dict)
-        for arg in checked_args:
-            check_type(scene[arg], annotations[arg])
-        print('Checked', path.stem)
+    scene_dict = check_type(check_scene_dict(safe_load(open(scene_path))), SceneDict)
+    scene = build_scene(scene_dict)
+    for arg in checked_args:
+        check_type(scene[arg], annotations[arg])
+    print('Checked', scene_path.stem)
 
 
-def test_rendering(scene_name: str) -> None:
-    scene_file = SCENES_PATH / (scene_name + '.yml')
-    scene = build_scene(check_scene_dict(safe_load(open(scene_file))))
-    images = check_type(render_scene(**scene, click=(-1, -1)), RenderedImages)
-    plot_rendered_images(images)
-
-
-def plot_rendered_images(images: RenderedImages) -> None:
-    image_names = images._fields
-    rows, cols = 3, 3
-    plt.style.use('grayscale')
-    _, axs = plt.subplots(rows, cols, figsize=(2 * cols, 2 * rows))
-    for ax, name in zip(axs.flatten(), image_names):
-        ax.imshow(to_rgb(getattr(images, name)))
-        ax.set_title(name.capitalize())
-    for ax in axs.flatten():
-        ax.axis('off')
-    plt.tight_layout()
-    plt.show()
+def test_rendering(scene_path: Path) -> Image:
+    scene = build_scene(check_scene_dict(safe_load(open(scene_path))))
+    t = time()
+    images = render_scene(**scene, click=(-1, -1))
+    print(f'Rendered {scene_path.stem:<15} | {time() - t:.2f} seconds')
+    images = check_type(images, RenderedImages)
+    return fromarray(hstack([to_rgb(image) for image in images]))
 
 
 if __name__ == '__main__':
-    if len(argv) == 1:
-        test_all_scenes()  # if no arguments, test build_scene on all scenes
-    elif len(argv) == 2:
-        test_rendering(argv[1])  # if one argument, test rendering that scene
-    else:
-        raise ValueError('Invalid number of arguments')
+    tmp = Path('tmp')
+    tmp.mkdir(exist_ok=True)
+    match tuple(argv[1:]):
+        case '-b',:  # build all scenes
+            for scene_path in SCENES_PATH.glob('*.yml'):
+                test_build_scene(scene_path)
+        case '-r',:  # render all scenes
+            for scene_path in SCENES_PATH.glob('*.yml'):
+                test_rendering(scene_path).save(tmp / f'{scene_path.stem}.png')
+        case '-r', scene_name:  # render a specific scene
+            test_rendering(SCENES_PATH / f'{scene_name}.yml').save(tmp / f'{scene_name}.png')
+        case _:
+            print('Usage: python tests.py [-b] [-r [scene_name]]')

utils/plot.py

@@ -31,6 +31,8 @@ def imshow(im: Optional[Array] = None, view_size: Tuple[int, int] = (0, 0)) -> g
 
 
 def to_rgb(im: Array) -> NpArray:
+    if len(im.shape) == 2:
+        im = im.reshape(*im.shape, 1).repeat(3, axis=2)
     if isnan(im).any():
         im = fill_nan(im)
     return uint8(255 * im.clip(0.0, 1.0))
@@ -49,7 +51,4 @@ def fill_nan(im: Array) -> Array:
     def color_nan_pixel(x: Array) -> Array:
         return np.where(np.isnan(x).any(), np.array([1, 0, 0]), x)
 
-    if len(im.shape) == 2:
-        im = np.tile(im.reshape(*im.shape, 1), (1, 1, 3))
-
     return vmap(vmap(color_nan_pixel))(im)