[GUI] Support images with 2 channels (RG) in GUI.set_image

docs/gui.rst

@@ -57,8 +57,7 @@ Paint on a window
 
     Set an image to display on the window.
 
-    The image pixels are set from the values of ``img[i, j]``, where ``i`` indicates the horizontal
-    coordinates (from left to right) and ``j`` the vertical coordinates (from bottom to top).
+    The image pixels are set from the values of ``img[i, j]``, where ``i`` indicates the horizontal coordinates (from left to right) and ``j`` the vertical coordinates (from bottom to top).
 
 
     If the window size is ``(x, y)``, then ``img`` must be one of:
@@ -67,12 +66,18 @@ Paint on a window
 
     * ``ti.field(shape=(x, y, 3))``, where `3` is for ``(r, g, b)`` channels
 
-    * ``ti.Vector.field(3, shape=(x, y))`` (see :ref:`vector`)
+    * ``ti.field(shape=(x, y, 2))``, where `2` is for ``(r, g)`` channels
+
+    * ``ti.Vector.field(3, shape=(x, y))`` ``(r, g, b)`` channels on each component (see :ref:`vector`)
+
+    * ``ti.Vector.field(2, shape=(x, y))`` ``(r, g)`` channels on each component
 
     * ``np.ndarray(shape=(x, y))``
 
     * ``np.ndarray(shape=(x, y, 3))``
 
+    * ``np.ndarray(shape=(x, y, 2))``
+
 
     The data type of ``img`` must be one of:
 
@@ -92,6 +97,13 @@ Paint on a window
         ``img`` entries will be clipped into range ``[0, 1]`` for display.
 
 
+.. function:: gui.get_image()
+
+    :return: (np.array) the current image shown on the GUI
+
+    Get the 4-channel (RGBA) image shown in the current GUI system.
+
+
 .. function:: gui.circle(pos, color = 0xFFFFFF, radius = 1)
 
     :parameter gui: (GUI) the window object
@@ -416,12 +428,6 @@ could make variable control more intuitive:
 Image I/O
 ---------
 
-.. function:: gui.get_image()
-
-    :return: a ``np.ndarray`` which is the current image shown on the GUI.
-
-    Get the RGBA shown image from the current GUI system which has four channels.
-
 .. function:: ti.imwrite(img, filename)
 
     :parameter img: (Matrix or Expr) the image you want to export

examples/stable_fluid.py

@@ -8,11 +8,11 @@
 import numpy as np
 import time
 
-res = 600
+res = 512             # 600 for a larger resoultion
 dt = 0.03
-p_jacobi_iters = 30
+p_jacobi_iters = 160  # 40 for quicker but not-so-accurate result
 f_strength = 10000.0
-curl_strength = 0
+curl_strength = 0     # 7 for unrealistic visual enhancement
 dye_decay = 0.99
 force_radius = res / 3.0
 debug = False
@@ -109,7 +109,7 @@ def backtrace_rk3(vf: ti.template(), p, dt: ti.template()):
     return p
 
 
-backtrace = backtrace_rk1
+backtrace = backtrace_rk3
 
 
 @ti.kernel
@@ -317,7 +317,10 @@ def reset():
         step(mouse_data)
 
     gui.set_image(dyes_pair.cur)
-    #gui.set_image(np.concatenate(
-    #        [np.abs(velocities_pair.cur.to_numpy() * 0.01 + 0.5),
-    #            np.zeros((res, res, 1))], axis=2))
+    # To visualize velocity field:
+    #gui.set_image(velocities_pair.cur.to_numpy() * 0.01 + 0.5)
+    # To visualize velocity divergence:
+    #divergence(velocities_pair.cur); gui.set_image(velocity_divs.to_numpy() * 0.1 + 0.5)
+    # To visualize velocity vorticity:
+    #vorticity(velocities_pair.cur); gui.set_image(velocity_curls.to_numpy() * 0.03 + 0.5)
     gui.show()

python/taichi/lang/meta.py

@@ -35,6 +35,8 @@ def vector_to_image(mat: ti.template(), arr: ti.ext_arr()):
     for I in ti.grouped(mat):
         for p in ti.static(range(mat.n)):
             arr[I, p] = cook_image_type(mat[I][p])
+            if ti.static(mat.n <= 2):
+                arr[I, 2] = 0
 
 
 @ti.kernel

python/taichi/misc/gui.py

@@ -107,6 +107,11 @@ def cook_image(self, img):
         if img.shape[2] == 3:
             zeros = np.zeros((img.shape[0], img.shape[1], 1), np.float32)
             img = np.concatenate([img, zeros], axis=2)
+        if img.shape[2] == 2:
+            zeros = np.zeros((img.shape[0], img.shape[1], 2), np.float32)
+            img = np.concatenate([img, zeros], axis=2)
+
+        assert img.shape[2] == 4, "Image must be grayscale, RG, RGB or RGBA"
 
         res = img.shape[:2]
         assert res == self.res, "Image resolution does not match GUI resolution"
@@ -138,12 +143,10 @@ def set_image(self, img):
                 self.img = self.cook_image(img.to_numpy())
             else:
                 # Type matched! We can use an optimized copy kernel.
-                assert img.shape \
-                 == self.res, "Image resolution does not match GUI resolution"
-                assert img.n in [
-                    3, 4
-                ], "Only greyscale, RGB or RGBA images are supported in GUI.set_image"
-                assert img.m == 1
+                assert img.shape  == self.res, \
+                        "Image resolution does not match GUI resolution"
+                assert img.n in [2, 3, 4] and img.m == 1, \
+                        "Only greyscale, RG, RGB or RGBA images are supported in GUI.set_image"
                 from taichi.lang.meta import vector_to_image
                 vector_to_image(img, self.img)
                 ti.sync()