[gui] Add GGUI set_image support for non-Vector fields and numpy ndarrays.

python/taichi/ui/staging_buffer.py

@@ -1,3 +1,4 @@
+import numpy as np
 from taichi.lang.impl import ndarray
 from taichi.lang.kernel_impl import kernel
 from taichi.lang.matrix import Vector
@@ -93,25 +94,81 @@ def copy_colors_to_vbo(vbo, colors):
 
 @ti.kernel
 def copy_image_f32_to_rgba8(src: ti.template(), dst: ti.template(),
-                            num_components: ti.template()):
-    for i, j in src:
+                            num_components: ti.template(),
+                            gray_scale: ti.template()):
+    for i, j in ti.ndrange(src.shape[0], src.shape[1]):
         px = ti.Vector([0, 0, 0, 0xff], dt=u32)
-        for k in ti.static(range(num_components)):
-            c = src[i, j][k]
+        if ti.static(gray_scale):
+            c = 0.0
+            c = src[i, j]
             c = max(0.0, min(1.0, c))
             c = c * 255
-            px[k] = ti.cast(c, u32)
+            px[0] = px[1] = px[2] = ti.cast(c, u32)
+        else:
+            for k in ti.static(range(num_components)):
+                c = 0.0
+                if ti.static(len(src.shape) == 3):
+                    c = src[i, j, k]
+                else:
+                    c = src[i, j][k]
+                c = max(0.0, min(1.0, c))
+                c = c * 255
+                px[k] = ti.cast(c, u32)
+        pack = (px[0] << 0 | px[1] << 8 | px[2] << 16 | px[3] << 24)
+        dst[i, j] = pack
+
+
+@ti.kernel
+def copy_image_f32_to_rgba8_np(src: ti.types.ndarray(), dst: ti.template(),
+                               num_components: ti.template(),
+                               gray_scale: ti.template()):
+    for i, j in ti.ndrange(src.shape[0], src.shape[1]):
+        px = ti.Vector([0, 0, 0, 0xff], dt=u32)
+        if ti.static(gray_scale):
+            c = 0.0
+            c = src[i, j]
+            c = max(0.0, min(1.0, c))
+            c = c * 255
+            px[0] = px[1] = px[2] = ti.cast(c, u32)
+        else:
+            for k in ti.static(range(num_components)):
+                c = src[i, j, k]
+                c = max(0.0, min(1.0, c))
+                c = c * 255
+                px[k] = ti.cast(c, u32)
         pack = (px[0] << 0 | px[1] << 8 | px[2] << 16 | px[3] << 24)
         dst[i, j] = pack
 
 
 @ti.kernel
 def copy_image_u8_to_rgba8(src: ti.template(), dst: ti.template(),
-                           num_components: ti.template()):
-    for i, j in src:
+                           num_components: ti.template(),
+                           gray_scale: ti.template()):
+    for i, j in ti.ndrange(src.shape[0], src.shape[1]):
         px = ti.Vector([0, 0, 0, 0xff], dt=u32)
-        for k in ti.static(range(num_components)):
-            px[k] = ti.cast(src[i, j][k], u32)
+        if ti.static(gray_scale):
+            px[0] = px[1] = px[2] = ti.cast(src[i, j], u32)
+        else:
+            for k in ti.static(range(num_components)):
+                if ti.static(len(src.shape) == 3):
+                    px[k] = ti.cast(src[i, j, k], u32)
+                else:
+                    px[k] = ti.cast(src[i, j][k], u32)
+        pack = (px[0] << 0 | px[1] << 8 | px[2] << 16 | px[3] << 24)
+        dst[i, j] = pack
+
+
+@ti.kernel
+def copy_image_u8_to_rgba8_np(src: ti.types.ndarray(), dst: ti.template(),
+                              num_components: ti.template(),
+                              gray_scale: ti.template()):
+    for i, j in ti.ndrange(src.shape[0], src.shape[1]):
+        px = ti.Vector([0, 0, 0, 0xff], dt=u32)
+        if ti.static(gray_scale):
+            px[0] = px[1] = px[2] = ti.cast(src[i, j], u32)
+        else:
+            for k in ti.static(range(num_components)):
+                px[k] = ti.cast(src[i, j, k], u32)
         pack = (px[0] << 0 | px[1] << 8 | px[2] << 16 | px[3] << 24)
         dst[i, j] = pack
 
@@ -122,24 +179,37 @@ def copy_image_u8_to_rgba8(src: ti.template(), dst: ti.template(),
 
 
 def to_rgba8(image):
-    if not hasattr(image, 'n') or image.m != 1:
-        raise Exception(
-            'the input image needs to be a Vector field (matrix with 1 column)'
-        )
-    if len(image.shape) != 2:
-        raise Exception(
-            "the shape of the image must be of the form (width,height)")
-
-    if image not in image_field_cache:
-        staging_img = ti.field(u32, image.shape)
-        image_field_cache[image] = staging_img
-    else:
-        staging_img = image_field_cache[image]
+    gray_scale = not hasattr(image, 'n') and len(image.shape) == 2
+    channels = 3
+    src_numpy = isinstance(image, np.ndarray)
+    if not gray_scale:
+        if len(image.shape) == 2:
+            channels = image.n
+        elif len(image.shape) == 3:
+            channels = image.shape[2]
+        else:
+            raise Exception(
+                "the shape of the image must be of the form (width,height) or (width,height,channels)"
+            )
 
-    if image.dtype == u8:
-        copy_image_u8_to_rgba8(image, staging_img, image.n)
-    elif image.dtype == f32:
-        copy_image_f32_to_rgba8(image, staging_img, image.n)
+    staging_key = image.shape[0:2] if src_numpy else image
+    if staging_key not in image_field_cache:
+        staging_img = ti.field(u32, image.shape[0:2])
+        image_field_cache[staging_key] = staging_img
+    else:
+        staging_img = image_field_cache[staging_key]
+
+    if image.dtype == u8 or image.dtype == np.uint8:
+        if src_numpy:
+            copy_image_u8_to_rgba8_np(image, staging_img, channels, gray_scale)
+        else:
+            copy_image_u8_to_rgba8(image, staging_img, channels, gray_scale)
+    elif image.dtype == f32 or image.dtype == np.float32:
+        if src_numpy:
+            copy_image_f32_to_rgba8_np(image, staging_img, channels,
+                                       gray_scale)
+        else:
+            copy_image_f32_to_rgba8(image, staging_img, channels, gray_scale)
     else:
         raise Exception("dtype of input image must either be u8 or f32")
     return staging_img