[GUI] Use Taichi kernels instead of NumPy operations to reduce GUI.set_image overhead

docs/gui.rst

@@ -89,7 +89,7 @@ Paint on a window
     .. note ::
 
         When using ``float32`` or ``float64`` as the data type,
-        ``img`` entries will be clipped into range ``[0, 1]``.
+        ``img`` entries will be clipped into range ``[0, 1]`` for display.
 
 
 .. function:: gui.circle(pos, color = 0xFFFFFF, radius = 1)
@@ -169,14 +169,14 @@ Every event have a key and type.
 
 ::
 
-  ti.GUI.ESCAPE
-  ti.GUI.SHIFT
-  ti.GUI.LEFT
-  'a'
+  ti.GUI.ESCAPE  # Esc
+  ti.GUI.SHIFT   # Shift
+  ti.GUI.LEFT    # Left Arrow
+  'a'            # we use lowercase for alphabet
   'b'
   ...
-  ti.GUI.LMB
-  ti.GUI.RMB
+  ti.GUI.LMB     # Left Mouse Button
+  ti.GUI.RMB     # Right Mouse Button
 
 *Event type* is the type of event, for now, there are just three type of event:
 

examples/simple_uv.py

@@ -0,0 +1,22 @@
+import taichi as ti
+import numpy as np
+
+ti.init()
+
+res = 1280, 720
+pixels = ti.Vector(3, dt=ti.f32, shape=res)
+
+
+@ti.kernel
+def paint():
+    for i, j in pixels:
+        u = i / res[0]
+        v = j / res[1]
+        pixels[i, j] = [u, v, 0]
+
+
+gui = ti.GUI('UV', res)
+while not gui.get_event(ti.GUI.ESCAPE):
+    paint()
+    gui.set_image(pixels)
+    gui.show()

python/taichi/lang/expr.py

@@ -146,6 +146,9 @@ def shape(self):
             self.materialize_layout_callback()
         return self.snode().shape()
 
+    def shape_ext(self):
+        return ()
+
     def data_type(self):
         return self.snode().data_type()
 

python/taichi/lang/matrix.py

@@ -604,6 +604,12 @@ def assign_renamed(x, y):
         from .meta import fill_matrix
         fill_matrix(self, val)
 
+    def shape_ext(self, as_vector=None):
+        if as_vector is None:
+            as_vector = self.m == 1
+        shape_ext = (self.n, ) if as_vector else (self.n, self.m)
+        return shape_ext
+
     @python_scope
     def to_numpy(self, keep_dims=False, as_vector=None):
         # Discussion: https://github.com/taichi-dev/taichi/pull/1046#issuecomment-633548858
@@ -615,8 +621,8 @@ def to_numpy(self, keep_dims=False, as_vector=None):
                 DeprecationWarning,
                 stacklevel=3)
         as_vector = self.m == 1 and not keep_dims
-        dim_ext = (self.n, ) if as_vector else (self.n, self.m)
-        ret = np.empty(self.loop_range().shape() + dim_ext,
+        shape_ext = self.shape_ext(as_vector)
+        ret = np.empty(self.loop_range().shape() + shape_ext,
                        dtype=to_numpy_type(
                            self.loop_range().snode().data_type()))
         from .meta import matrix_to_ext_arr
@@ -629,8 +635,8 @@ def to_numpy(self, keep_dims=False, as_vector=None):
     def to_torch(self, device=None, keep_dims=False):
         import torch
         as_vector = self.m == 1 and not keep_dims
-        dim_ext = (self.n, ) if as_vector else (self.n, self.m)
-        ret = torch.empty(self.loop_range().shape() + dim_ext,
+        shape_ext = self.shape_ext(as_vector)
+        ret = torch.empty(self.loop_range().shape() + shape_ext,
                           dtype=to_pytorch_type(
                               self.loop_range().snode().data_type()),
                           device=device)

python/taichi/lang/meta.py

@@ -15,6 +15,28 @@ def tensor_to_ext_arr(tensor: ti.template(), arr: ti.ext_arr()):
         arr[I] = tensor[I]
 
 
+@ti.func
+def cook_image_type(x):
+    x = ti.cast(x, ti.f32)
+    # Issue #1000, don't know why win GUI doesn't do clamp for us
+    if ti.static(ti.get_os_name() == 'win'):
+        x = ti.min(1, ti.max(0, x))
+    return x
+
+
+@ti.kernel
+def tensor_to_image(tensor: ti.template(), arr: ti.ext_arr()):
+    for I in ti.grouped(tensor):
+        arr[I] = cook_image_type(tensor[I])
+
+
+@ti.kernel
+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])
+
+
 @ti.kernel
 def tensor_to_tensor(tensor: ti.template(), other: ti.template()):
     for I in ti.grouped(tensor):

python/taichi/misc/gui.py

@@ -32,6 +32,8 @@ def __init__(self, name, res=512, background_color=0x0):
         if isinstance(res, numbers.Number):
             res = (res, res)
         self.res = res
+        # The GUI canvas uses RGBA for storage, therefore we needs NxMx4 for a image
+        self.img = np.ascontiguousarray(np.zeros(self.res + (4, ), np.float32))
         self.core = ti.core.GUI(name, ti.veci(*res))
         self.canvas = self.core.get_canvas()
         self.background_color = background_color
@@ -49,30 +51,69 @@ def clear(self, color=None):
 
     def set_image(self, img):
         import numpy as np
-        from .image import cook_image
-        img = cook_image(img)
-        if img.dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
-            img = img.astype(np.float32) * (1 / np.iinfo(img.dtype).max)
-        elif img.dtype in [np.float32, np.float64]:
-            img = np.clip(img.astype(np.float32), 0, 1)
+        import taichi as ti
+
+        def cook_image(img):
+            if img.dtype in [np.uint8, np.uint16, np.uint32, np.uint64]:
+                img = img.astype(np.float32) * (1 / np.iinfo(img.dtype).max)
+            elif img.dtype in [np.float32, np.float64]:
+                img = img.astype(np.float32)
+                from .util import get_os_name
+                if img.dtype == np.float32 and get_os_name() == 'win':
+                    img = np.clip(img, 0, 1)
+            else:
+                raise ValueError(
+                    f'Data type {img.dtype} not supported in GUI.set_image')
+
+            if len(img.shape) == 2:
+                img = img[..., None]
+            if img.shape[2] == 1:
+                img = img + np.zeros(shape=(1, 1, 4), dtype=np.float32)
+            if img.shape[2] == 3:
+                img = np.concatenate([
+                    img,
+                    np.zeros(shape=(img.shape[0], img.shape[1], 1),
+                             dtype=np.float32)
+                ],
+                                     axis=2)
+            assert img.shape[:
+                             2] == self.res, "Image resolution does not match GUI resolution"
+            return np.ascontiguousarray(img)
+
+        if isinstance(img, ti.Expr):
+            if ti.core.is_integral(img.data_type()):
+                # image of uint is not optimized by xxx_to_image
+                self.img = cook_image(img.to_numpy())
+            else:
+                assert img.shape(
+                ) == self.res, "Image resolution does not match GUI resolution"
+                from taichi.lang.meta import tensor_to_image
+                tensor_to_image(img, self.img)
+                ti.sync()
+
+        elif isinstance(img, ti.Matrix):
+            if ti.core.is_integral(img.data_type()):
+                self.img = cook_image(img.to_numpy())
+            else:
+                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
+                from taichi.lang.meta import vector_to_image
+                vector_to_image(img, self.img)
+                ti.sync()
+
+        elif isinstance(img, np.ndarray):
+            self.img = cook_image(img)
+
         else:
             raise ValueError(
-                f'Data type {img.dtype} not supported in GUI.set_image')
-        if len(img.shape) == 2:
-            img = img[..., None]
-        if img.shape[2] == 1:
-            img = img + np.zeros(shape=(1, 1, 4))
-        if img.shape[2] == 3:
-            img = np.concatenate([
-                img,
-                np.zeros(shape=(img.shape[0], img.shape[1], 1),
-                         dtype=np.float32)
-            ],
-                                 axis=2)
-        img = img.astype(np.float32)
-        assert img.shape[:
-                         2] == self.res, "Image resolution does not match GUI resolution"
-        self.core.set_img(np.ascontiguousarray(img).ctypes.data)
+                f"GUI.set_image only takes Taichi tensor or NumPy array, not {type(img)}"
+            )
+
+        self.core.set_img(self.img.ctypes.data)
 
     def circle(self, pos, color=0xFFFFFF, radius=1):
         self.canvas.circle_single(pos[0], pos[1], color, radius)

python/taichi/misc/image.py

@@ -2,18 +2,10 @@
 import taichi as ti
 
 
-def cook_image(img):
-    if isinstance(img, ti.Matrix):
-        img = img.to_numpy(as_vector=True)
-    if isinstance(img, ti.Expr):
+def imwrite(img, filename):
+    if not isinstance(img, np.ndarray):
         img = img.to_numpy()
-    assert isinstance(img, np.ndarray)
     assert len(img.shape) in [2, 3]
-    return img
-
-
-def imwrite(img, filename):
-    img = cook_image(img)
     resx, resy = img.shape[:2]
     if len(img.shape) == 3:
         comp = img.shape[2]
@@ -35,7 +27,9 @@ def imread(filename, channels=0):
 
 
 def imshow(img, winname='Taichi'):
-    img = cook_image(img)
+    if not isinstance(img, np.ndarray):
+        img = img.to_numpy()
+    assert len(img.shape) in [2, 3]
     gui = ti.GUI(winname, res=img.shape[:2])
     while not gui.get_event(ti.GUI.ESCAPE):
         gui.set_image(img)