[GUI] Support saving image without GUI showing a window

python/taichi/misc/gui.py

@@ -34,14 +34,21 @@ class Event:
     PRESS = ti_core.KeyEvent.EType.Press
     RELEASE = ti_core.KeyEvent.EType.Release
 
-    def __init__(self, name='Taichi', res=512, background_color=0x0):
+    def __init__(self,
+                 name='Taichi',
+                 show_GUI=True,
+                 res=512,
+                 background_color=0x0):
         self.name = name
         if isinstance(res, numbers.Number):
             res = (res, res)
         self.res = res
         # The GUI canvas uses RGBA for storage, therefore we need NxMx4 for an image.
         self.img = np.ascontiguousarray(np.zeros(self.res + (4, ), np.float32))
+        self.show_GUI = show_GUI
         self.core = ti_core.GUI(name, core_veci(*res))
+        if self.show_GUI:
+            self.core.initialise_window()
         self.canvas = self.core.get_canvas()
         self.background_color = background_color
         self.key_pressed = set()
@@ -356,7 +363,8 @@ def arrow_field(self, dir, radius=1, color=0xffffff, bound=0.5, **kwargs):
         self.arrows(base, dir, radius=radius, color=color, **kwargs)
 
     def show(self, file=None):
-        self.core.update()
+        if self.show_GUI:
+            self.core.update()
         if file:
             self.core.screenshot(file)
         self.frame += 1

taichi/gui/gui.h

@@ -431,7 +431,7 @@ class Canvas {
   ~Canvas() {
   }
 
-  void set_idendity_transform_matrix() {
+  void set_identity_transform_matrix() {
     transform_matrix = Matrix3(1);
   }
 };
@@ -769,10 +769,8 @@ class GUI : public GUIBase {
     buffer.initialize(Vector2i(width, height));
     canvas = std::make_unique<Canvas>(buffer);
     last_frame_time = taichi::Time::get_time();
-    create_window();
-    set_title(window_name);
     if (!normalized_coord) {
-      canvas->set_idendity_transform_matrix();
+      canvas->set_identity_transform_matrix();
     }
     widget_height = 0;
   }
@@ -785,6 +783,11 @@ class GUI : public GUIBase {
 
   void create_window();
 
+  void initialise_window() {
+    create_window();
+    set_title(window_name);
+  }
+
   Canvas &get_canvas() {
     return *canvas;
   }
@@ -795,7 +798,7 @@ class GUI : public GUIBase {
 
   void redraw_widgets() {
     auto old_transform_matrix = canvas->transform_matrix;
-    canvas->set_idendity_transform_matrix();
+    canvas->set_identity_transform_matrix();
     for (auto &w : widgets) {
       w->set_hover(w->inside(cursor_pos));
       w->redraw(*canvas);

taichi/python/export_visual.cpp

@@ -29,6 +29,7 @@ void export_visual(py::module &m) {
       .def_readwrite("frame_delta_limit", &GUI::frame_delta_limit)
       .def_readwrite("should_close", &GUI::should_close)
       .def("get_canvas", &GUI::get_canvas, py::return_value_policy::reference)
+      .def("initialise_window", [](GUI *gui) { gui->initialise_window(); })
       .def("set_img",
            [&](GUI *gui, std::size_t ptr) {
              auto &img = gui->canvas->img;

tests/python/test_gui.py

@@ -0,0 +1,31 @@
+from taichi import make_temp_file
+import taichi as ti
+import numpy as np
+from PIL import Image
+import pytest
+
+
+@ti.host_arch_only
+@pytest.mark.parametrize('dtype', [ti.u8, ti.f32])
+def test_save_image_without_window(dtype):
+    n = 255
+    pixels = ti.field(dtype=dtype, shape=(n, n, 3))
+
+    @ti.kernel
+    def paint(c: dtype):
+        for i, j, k in pixels:
+            pixels[i, j, k] = c
+
+    gui = ti.GUI("Test", show_GUI=False, res=(n, n))
+    for i in [0, 32, 64, 128, 255]:
+        if dtype is ti.u8:
+            paint(i)
+        else:
+            paint(i * 1.0 / n)
+        gui.set_image(pixels)
+        image_path = make_temp_file(suffix='.jpg')
+        gui.show(image_path)
+        image = np.array(Image.open(image_path))
+        delta = (image - i).sum()
+        assert delta == 0, "Expected image difference to be 0 but got {} instead.".format(
+            delta)