diff --git a/docs/gui.rst b/docs/gui.rst index 94a50322d20b2..436db4dce7116 100644 --- a/docs/gui.rst +++ b/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: diff --git a/examples/image_io.py b/examples/image_io.py deleted file mode 100644 index 11b8aa440f01d..0000000000000 --- a/examples/image_io.py +++ /dev/null @@ -1,24 +0,0 @@ -import taichi as ti -import os - -pixel = ti.var(ti.u8, shape=(512, 512, 3)) - - -@ti.kernel -def paint(): - for I in ti.grouped(pixel): - pixel[I] = ti.random() * 255 - - -paint() -pixel = pixel.to_numpy() -ti.imshow(pixel, 'Random Generated') -for ext in ['bmp', 'png', 'jpg']: - fn = 'taichi-example-random-img.' + ext - ti.imwrite(pixel, fn) - pixel_r = ti.imread(fn) - if ext != 'jpg': - assert (pixel_r == pixel).all() - else: - ti.imshow(pixel_r, 'JPEG Read Result') - os.remove(fn) diff --git a/examples/simple_uv.py b/examples/simple_uv.py new file mode 100644 index 0000000000000..f9e30568c1ede --- /dev/null +++ b/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() diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py index 80229a3d14c61..45c212b6e70c0 100644 --- a/python/taichi/lang/matrix.py +++ b/python/taichi/lang/matrix.py @@ -3,7 +3,7 @@ import copy import numbers import numpy as np -from .util import taichi_scope, python_scope, deprecated, to_numpy_type +from .util import taichi_scope, python_scope, deprecated, to_numpy_type, to_pytorch_type from .common_ops import TaichiOperations from collections.abc import Iterable @@ -631,12 +631,12 @@ 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) + shape_ext = (self.n, ) if as_vector else (self.n, self.m) if self.is_pyconstant(): - return np.array(self.entries).reshape(dim_ext) + return np.array(self.entries).reshape(shape_ext) - ret = np.empty(self.loop_range().shape() + dim_ext, + 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 @@ -649,8 +649,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.n, ) if as_vector else (self.n, self.m) + ret = torch.empty(self.loop_range().shape() + shape_ext, dtype=to_pytorch_type( self.loop_range().snode().data_type()), device=device) diff --git a/python/taichi/lang/meta.py b/python/taichi/lang/meta.py index 9fdc8e231461c..ae38f5f7b5a5f 100644 --- a/python/taichi/lang/meta.py +++ b/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) + return x + + +@ti.kernel +def tensor_to_image(tensor: ti.template(), arr: ti.ext_arr()): + for I in ti.grouped(tensor): + t = cook_image_type(tensor[I]) + arr[I, 0] = t + arr[I, 1] = t + arr[I, 2] = t + + +@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): diff --git a/python/taichi/misc/gui.py b/python/taichi/misc/gui.py index e4e36918de2e9..0205ccd651a26 100644 --- a/python/taichi/misc/gui.py +++ b/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 need NxMx4 for an 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 @@ -47,32 +49,66 @@ def clear(self, color=None): color = self.background_color self.canvas.clear(color) - def set_image(self, img): - import numpy as np - from .image import cook_image - img = cook_image(img) + def cook_image(self, 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) + img = img.astype(np.float32) 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)) + img = img + np.zeros((1, 1, 4), 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) - 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) + zeros = np.zeros((img.shape[0], img.shape[1], 1), np.float32) + img = np.concatenate([img, zeros], axis=2) + + res = img.shape[:2] + assert res == self.res, "Image resolution does not match GUI resolution" + return np.ascontiguousarray(img) + + def set_image(self, img): + import numpy as np + import taichi as ti + + 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 = self.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 = self.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 = self.cook_image(img) + + else: + raise ValueError( + f"GUI.set_image only takes a 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) diff --git a/python/taichi/misc/image.py b/python/taichi/misc/image.py index 30469709e4820..09cbb1fd98973 100644 --- a/python/taichi/misc/image.py +++ b/python/taichi/misc/image.py @@ -2,23 +2,28 @@ 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 + if img.dtype in [np.uint16, np.uint32, np.uint64]: + img = (img // (np.iinfo(img.dtype).max / 256)).astype(np.uint8) + elif img.dtype in [np.float32, np.float64]: + img = (np.clip(img, 0, 1) * 255.0 + 0.5).astype(np.uint8) + elif img.dtype != np.uint8: + raise ValueError(f'Data type {img.dtype} not supported in ti.imwrite') + + assert len(img.shape) in [2, + 3], "Image must be either RGB/RGBA or greyscale" + assert img.shape[2] in [1, 3, + 4], "Image must be either RGB/RGBA or greyscale" -def imwrite(img, filename): - img = cook_image(img) resx, resy = img.shape[:2] - if len(img.shape) == 3: - comp = img.shape[2] - else: + if len(img.shape) == 2: comp = 1 + else: + comp = img.shape[2] + img = np.ascontiguousarray(img.swapaxes(0, 1)[::-1, :, :]) ptr = img.ctypes.data ti.core.imwrite(filename, ptr, resx, resy, comp) @@ -34,9 +39,13 @@ def imread(filename, channels=0): return img.swapaxes(0, 1)[:, ::-1, :] -def imshow(img, winname='Taichi'): - img = cook_image(img) - gui = ti.GUI(winname, res=img.shape[:2]) +def imshow(img, window_name='Taichi'): + if not isinstance(img, np.ndarray): + img = img.to_numpy() + assert len(img.shape) in [2, + 3], "Image must be either RGB/RGBA or greyscale" + gui = ti.GUI(window_name, res=img.shape[:2]) + img = gui.cook_image(img) while not gui.get_event(ti.GUI.ESCAPE): gui.set_image(img) gui.show() diff --git a/taichi/gui/win32.cpp b/taichi/gui/win32.cpp index eb3e045c1de0b..0c24ead17c5fc 100644 --- a/taichi/gui/win32.cpp +++ b/taichi/gui/win32.cpp @@ -180,9 +180,10 @@ void GUI::redraw() { for (int i = 0; i < width; i++) { for (int j = 0; j < height; j++) { auto c = reinterpret_cast(data + (j * width) + i); - c[0] = (unsigned char)(canvas->img[i][height - j - 1][2] * 255.0_f); - c[1] = (unsigned char)(canvas->img[i][height - j - 1][1] * 255.0_f); - c[2] = (unsigned char)(canvas->img[i][height - j - 1][0] * 255.0_f); + auto d = canvas->img[i][height - j - 1]; + c[0] = uint8(clamp(int(d[2] * 255.0_f), 0, 255)); + c[1] = uint8(clamp(int(d[1] * 255.0_f), 0, 255)); + c[2] = uint8(clamp(int(d[0] * 255.0_f), 0, 255)); c[3] = 0; } } diff --git a/taichi/gui/x11.cpp b/taichi/gui/x11.cpp index 84aaf919ef174..b8f2e0a02dc69 100644 --- a/taichi/gui/x11.cpp +++ b/taichi/gui/x11.cpp @@ -35,7 +35,7 @@ class CXImage { *p++ = uint8(clamp(int(c[2] * 255.0_f), 0, 255)); *p++ = uint8(clamp(int(c[1] * 255.0_f), 0, 255)); *p++ = uint8(clamp(int(c[0] * 255.0_f), 0, 255)); - *p++ = uint8(clamp(int(c[3] * 255.0_f), 0, 255)); + *p++ = 0; } } }