[GUI] Add handy ti.imscale(img, w, h) for scaling image (#1735)

Co-authored-by: Xudong Feng <[email protected]>

docs/gui.rst

@@ -514,6 +514,7 @@ Image I/O
 
     Each value in this returned field is an integer in [0, 255].
 
+
 .. function:: ti.imshow(img, windname)
 
     :parameter img: (ti.Vector.field or ti.field) the image to show in the GUI
@@ -522,3 +523,15 @@ Image I/O
     This function will create an instance of ``ti.GUI`` and show the input image on the screen.
 
     It has the same logic as ``ti.imwrite`` for different datatypes.
+
+
+.. function:: ti.imresize(img, w, h=None):
+
+    :parameter img: (np.array or ti.field) the input image.
+    :parameter w: (int) the width after resizing.
+    :parameter h: (optional, int) the height after resizing.
+    :return: (np.array) the resized image.
+
+    If ``h`` is not specified, it will be equal to ``w`` by default.
+
+    The output image shape is: ``(w, h, *img.shape[2:])``.

examples/game_of_life.py

@@ -11,7 +11,6 @@
 img_size = n * cell_size
 alive = ti.field(int, shape=(n, n))  # alive = 1, dead = 0
 count = ti.field(int, shape=(n, n))  # count of neighbours
-img = ti.field(float, shape=(img_size, img_size))  # image to be displayed
 
 
 @ti.func
@@ -58,13 +57,6 @@ def init():
             alive[i, j] = 0
 
 
-@ti.kernel
-def render():
-    for i, j in alive:
-        for u, v in ti.static(ti.ndrange(cell_size, cell_size)):
-            img[i * cell_size + u, j * cell_size + v] = alive[i, j]
-
-
 gui = ti.GUI('Game of Life', (img_size, img_size))
 gui.fps_limit = 15
 
@@ -91,7 +83,5 @@ def render():
     if not paused:
         run()
 
-    render()
-
-    gui.set_image(img)
+    gui.set_image(ti.imresize(alive, img_size).astype(np.uint8) * 255)
     gui.show()

python/taichi/misc/image.py

@@ -34,7 +34,11 @@ def imdisplay(img):
     """
     Try to display image in interactive shell.
     """
-    if ti.lang.shell.oinspect.name == ti.lang.shell.ShellType.JUPYTER:
+    try:
+        get_ipython()
+    except:
+        ti.imshow(img)
+    else:
         import PIL.Image
         from io import BytesIO
         import IPython.display
@@ -43,8 +47,21 @@ def imdisplay(img):
         with BytesIO() as f:
             PIL.Image.fromarray(img).save(f, 'png')
             IPython.display.display(IPython.display.Image(data=f.getvalue()))
-    else:
-        ti.imshow(img)
+
+
+def imresize(img, w, h=None):
+    """
+    Resize an image to a specific size.
+    """
+    if not isinstance(img, np.ndarray):
+        img = img.to_numpy()
+    if h is None:
+        h = w
+    assert isinstance(w, int) and isinstance(h, int) and w > 1 and h > 1
+    u, v = (img.shape[0]) / (w), (img.shape[1]) / (h)
+    x = np.clip(np.arange(w) * u, 0, img.shape[0] - 1).astype(np.int32)
+    y = np.clip(np.arange(h) * v, 0, img.shape[1] - 1).astype(np.int32)
+    return img[tuple(np.meshgrid(x, y))].swapaxes(0, 1)
 
 
 def imwrite(img, filename):
@@ -94,5 +111,6 @@ def imshow(img, window_name='imshow'):
     'imshow',
     'imread',
     'imwrite',
+    'imresize',
     'imdisplay',
 ]

tests/python/test_image_io.py

@@ -69,3 +69,19 @@ def test_image_io_uint(resx, resy, comp, ext, dt):
     pixel_r = ti.imread(fn).astype(np_type) * np_max
     assert (pixel_r == pixel).all()
     os.remove(fn)
+
+
+@pytest.mark.parametrize('comp', [1, 3])
+@pytest.mark.parametrize('resx,resy', [(91, 81)])
+@pytest.mark.parametrize('scale', [1, 2, 3])
+@ti.host_arch_only
+def test_image_resize_sum(resx, resy, comp, scale):
+    shape = (resx, resy)
+    if comp != 1:
+        shape = shape + (comp, )
+    old_img = np.random.rand(*shape).astype(np.float32)
+    if resx == resy:
+        new_img = ti.imresize(old_img, resx * scale)
+    else:
+        new_img = ti.imresize(old_img, resx * scale, resy * scale)
+    assert np.sum(old_img) * scale**2 == ti.approx(np.sum(new_img))