[test] [gui] Add a test of fetching color attachment

tests/python/test_ggui.py

@@ -16,39 +16,27 @@
 supported_archs = [ti.vulkan, ti.cuda]
 
 
-def get_temp_png():
-    f, name = tempfile.mkstemp(suffix='.png')
-    os.close(f)
-    return name
-
-
-def write_temp_image(window):
-    f = get_temp_png()
-    window.save_image(f)
-    try:
-        os.remove(f)
-    except OSError:
-        pass
-
-
-def verify_image(window, image_name, tolerance=0.1):
+def verify_image(image, image_name, tolerance=0.1):
     if REGENERATE_GROUNDTRUTH_IMAGES:
         ground_truth_name = f"tests/python/expected/{image_name}.png"
-        window.save_image(ground_truth_name)
+        ti.tools.imwrite(image, ground_truth_name)
     else:
         ground_truth_name = str(
             pathlib.Path(__file__).parent) + f"/expected/{image_name}.png"
-        actual_name = get_temp_png()
-        window.save_image(actual_name)
         ground_truth_np = ti.tools.imread(ground_truth_name)
-        actual_np = ti.tools.imread(actual_name)
-        assert len(ground_truth_np.shape) == len(actual_np.shape)
-        for i in range(len(ground_truth_np.shape)):
-            assert ground_truth_np.shape[i] == actual_np.shape[i]
-        diff = ground_truth_np - actual_np
-        mse = np.mean(diff * diff)
-        assert mse <= tolerance  # the pixel values are 0~255
-        os.remove(actual_name)
+
+        with tempfile.NamedTemporaryFile(suffix='.png') as fp:
+            actual_name = fp.name
+            ti.tools.imwrite(image, actual_name)
+            actual_np = ti.tools.imread(actual_name)
+
+            assert len(ground_truth_np.shape) == len(actual_np.shape)
+            for i in range(len(ground_truth_np.shape)):
+                assert ground_truth_np.shape[i] == actual_np.shape[i]
+
+            diff = ground_truth_np - actual_np
+            mse = np.mean(diff * diff)
+            assert mse <= tolerance  # the pixel values are 0~255
 
 
 @pytest.mark.skipif(not _ti_core.GGUI_AVAILABLE, reason="GGUI Not Available")
@@ -142,13 +130,15 @@ def render():
 
     for _ in range(RENDER_REPEAT):
         render()
-        write_temp_image(window)
+        window.get_image_buffer_as_numpy()
+
     render()
     if (platform.system() == 'Darwin'):
         # FIXME: Use lower tolerance when macOS ggui supports wide lines
-        verify_image(window, 'test_geometry_2d', 1.0)
+        verify_image(window.get_image_buffer_as_numpy(), 'test_geometry_2d',
+                     1.0)
     else:
-        verify_image(window, 'test_geometry_2d')
+        verify_image(window.get_image_buffer_as_numpy(), 'test_geometry_2d')
     window.destroy()
 
 
@@ -238,9 +228,9 @@ def render():
 
     for _ in range(RENDER_REPEAT):
         render()
-        write_temp_image(window)
+        window.get_image_buffer_as_numpy()
     render()
-    verify_image(window, 'test_geometry_3d')
+    verify_image(window.get_image_buffer_as_numpy(), 'test_geometry_3d')
     window.destroy()
 
 
@@ -264,9 +254,9 @@ def render():
 
     for _ in range(RENDER_REPEAT):
         render()
-        write_temp_image(window)
+        window.get_image_buffer_as_numpy()
     render()
-    verify_image(window, 'test_set_image')
+    verify_image(window.get_image_buffer_as_numpy(), 'test_set_image')
     window.destroy()
 
 
@@ -293,16 +283,16 @@ def render():
 
     for _ in range(RENDER_REPEAT):
         render()
-        write_temp_image(window)
+        window.get_image_buffer_as_numpy()
     render()
-    verify_image(window, 'test_set_image')
+    verify_image(window.get_image_buffer_as_numpy(), 'test_set_image')
     window.destroy()
 
 
 @pytest.mark.skipif(not _ti_core.GGUI_AVAILABLE, reason="GGUI Not Available")
 @test_utils.test(arch=[ti.vulkan])
 def test_set_image_with_texture():
-    window = ti.ui.Window('test', (640, 480), show_window=True)
+    window = ti.ui.Window('test', (640, 480), show_window=False)
     canvas = window.get_canvas()
 
     img = ti.Texture(ti.f32, 4, (512, 512))
@@ -323,9 +313,9 @@ def render():
 
     for _ in range(3):
         render()
-        write_temp_image(window)
+        window.get_image_buffer_as_numpy()
     render()
-    verify_image(window, 'test_set_image')
+    verify_image(window.get_image_buffer_as_numpy(), 'test_set_image')
     window.destroy()
 
 
@@ -347,9 +337,9 @@ def render():
 
     for _ in range(RENDER_REPEAT):
         render()
-        write_temp_image(window)
+        window.get_image_buffer_as_numpy()
     render()
-    verify_image(window, 'test_imgui')
+    verify_image(window.get_image_buffer_as_numpy(), 'test_imgui')
     window.destroy()
 
 
@@ -381,3 +371,30 @@ def test_get_camera_view_and_projection_matrix():
     assert (abs(projection_matrix[2, 2] - 1.0001000e-4) <= 1e-5)
     assert (abs(projection_matrix[2, 3] + 1.0000000) <= 1e-5)
     assert (abs(projection_matrix[3, 2] - 1.0001000e-1) <= 1e-5)
+
+
+@pytest.mark.skipif(not _ti_core.GGUI_AVAILABLE, reason="GGUI Not Available")
+@test_utils.test(arch=supported_archs)
+def test_fetching_color_attachment():
+    window = ti.ui.Window('test', (640, 480), show_window=False)
+    canvas = window.get_canvas()
+
+    img = ti.Vector.field(4, ti.f32, (512, 512))
+
+    @ti.kernel
+    def init_img():
+        for i, j in img:
+            img[i, j] = ti.Vector([i, j, 0, 512], dt=ti.f32) / 512
+
+    init_img()
+
+    def render():
+        canvas.set_image(img)
+
+    for _ in range(RENDER_REPEAT):
+        render()
+        window.get_image_buffer_as_numpy()
+
+    render()
+    verify_image(window.get_image_buffer_as_numpy(), 'test_set_image')
+    window.destroy()