[Bug] [misc] Fix image I/O for channels = 1 and improve test coverage…

… (unrevert) (#1242)

README.md

@@ -58,7 +58,7 @@ Please build from source for other configurations (e.g., your CPU is ARM).
 - [LBM Taichi](https://github.com/hietwll/LBM_Taichi): Fluid solver based on Lattice Boltzmann method implemented by Taichi programming language, by [Zhuo Wang (hietwll)](https://github.com/hietwll).
 - [Shadertoy reproduced by Taichi](https://github.com/Phonicavi/Shadertoy-taichi): Some prevalent shadertoys implemented in Taichi, by [QIU Feng (Phonicavi)](https://github.com/Phonicavi).
 - [DiffTaichi](https://github.com/yuanming-hu/difftaichi): 10 differentiable physical simulators built with Taichi differentiable programming, by [Yuanming Hu (yuanming-hu)](https://github.com/yuanming-hu).
-- [Taichi GLSL](https://github.com/yuanming-hu/difftaichi): Manipulate Taichi with GLSL-alike helper functions.
+- [Taichi GLSL](https://github.com/taichi-dev/taichi_glsl): Manipulate Taichi with GLSL-alike helper functions.
 
 ## Developers
 

python/build.py

@@ -86,10 +86,19 @@ def get_python_executable():
     shutil.copy('../runtimes/RelWithDebInfo/taichi_core.dll',
                 'taichi/lib/taichi_core.pyd')
 
-os.system('{} -m taichi changelog --save && cat ../CHANGELOG.md'.format(
-    get_python_executable()))
+os.system(f'cd .. && {get_python_executable()} -m taichi changelog --save')
+
+try:
+    with open('../CHANGELOG.md') as f:
+        print(f.read())
+except FileNotFoundError:
+    print('CHANGELOG.md not found')
+    pass
 
-shutil.copy('../CHANGELOG.md', './taichi/CHANGELOG.md')
+try:
+    shutil.copy('../CHANGELOG.md', './taichi/CHANGELOG.md')
+except FileNotFoundError:
+    pass
 shutil.copytree('../tests/python', './taichi/tests')
 shutil.copytree('../examples', './taichi/examples')
 shutil.copytree('../external/assets', './taichi/assets')
@@ -134,9 +143,8 @@ def get_python_executable():
             '%PYPI_PWD%' if get_os_name() == 'win' else '$PYPI_PWD'))
 elif mode == 'test':
     print('Uninstalling old taichi packages...')
-    os.system('{} -m pip uninstall taichi-nightly'.format(
-        get_python_executable()))
-    os.system('{} -m pip uninstall taichi'.format(get_python_executable()))
+    os.system(f'{get_python_executable()} -m pip uninstall taichi-nightly')
+    os.system(f'{get_python_executable()} -m pip uninstall taichi')
     dists = os.listdir('dist')
     assert len(dists) == 1
     dist = dists[0]

python/taichi/misc/image.py

@@ -7,24 +7,23 @@ def imwrite(img, filename):
         img = img.to_numpy()
 
     if img.dtype in [np.uint16, np.uint32, np.uint64]:
-        img = (img // (np.iinfo(img.dtype).max / 256)).astype(np.uint8)
+        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"
 
     resx, resy = img.shape[:2]
     if len(img.shape) == 2:
         comp = 1
     else:
         comp = img.shape[2]
+    assert comp in [1, 3, 4], "Image must be either RGB/RGBA or greyscale"
 
-    img = np.ascontiguousarray(img.swapaxes(0, 1)[::-1, :, :])
+    img = np.ascontiguousarray(img.swapaxes(0, 1)[::-1, :])
     ptr = img.ctypes.data
     ti.core.imwrite(filename, ptr, resx, resy, comp)
 

tests/python/test_image_io.py

@@ -1,16 +1,77 @@
 import taichi as ti
 import numpy as np
+import pytest
 import os
+from tempfile import mkstemp
 
 
+def make_temp(*args, **kwargs):
+    fd, name = mkstemp(*args, **kwargs)
+    os.close(fd)
+    return name
+
+
+# jpg is also supported but hard to test here since it's lossy:
+@pytest.mark.parametrize('comp,ext', [(3, 'bmp'), (1, 'png'), (3, 'png'),
+                                      (4, 'png')])
+@pytest.mark.parametrize('resx,resy', [(201, 173)])
+@pytest.mark.parametrize('is_tensor', [False, True])
+@pytest.mark.parametrize('dt', [ti.u8])
 @ti.host_arch_only
-def test_image_io():
-    pixel = (np.random.rand(201, 173, 3) * 255).astype(np.uint8)
-    for ext in [
-            'bmp', 'png'
-    ]:  # jpg is also supported but hard to test here since it's lossy
-        fn = 'taichi-image-io-test.' + ext
+def test_image_io(resx, resy, comp, ext, is_tensor, dt):
+    if comp != 1:
+        shape = (resx, resy, comp)
+    else:
+        shape = (resx, resy)
+    if is_tensor:
+        pixel_t = ti.var(dt, shape)
+    pixel = np.random.randint(256, size=shape, dtype=ti.to_numpy_type(dt))
+    if is_tensor:
+        pixel_t.from_numpy(pixel)
+    fn = make_temp(suffix='.' + ext)
+    if is_tensor:
+        ti.imwrite(pixel_t, fn)
+    else:
         ti.imwrite(pixel, fn)
-        pixel_r = ti.imread(fn)
-        assert (pixel_r == pixel).all()
-        os.remove(fn)
+    pixel_r = ti.imread(fn)
+    if comp == 1:
+        # from (resx, resy, 1) to (resx, resy)
+        pixel_r = pixel_r.reshape((resx, resy))
+    assert (pixel_r == pixel).all()
+    os.remove(fn)
+
+
+@pytest.mark.parametrize('comp,ext', [(3, 'png'), (4, 'png')])
+@pytest.mark.parametrize('resx,resy', [(91, 81)])
+@pytest.mark.parametrize('dt', [ti.f32, ti.f64])
+@ti.host_arch_only
+def test_image_io_vector(resx, resy, comp, ext, dt):
+    shape = (resx, resy)
+    pixel = np.random.rand(*shape, comp).astype(ti.to_numpy_type(dt))
+    pixel_t = ti.Vector(comp, dt, shape)
+    pixel_t.from_numpy(pixel)
+    fn = make_temp(suffix='.' + ext)
+    ti.imwrite(pixel_t, fn)
+    pixel_r = (ti.imread(fn).astype(ti.to_numpy_type(dt)) + 0.5) / 256.0
+    assert np.allclose(pixel_r, pixel, atol=2e-2)
+    os.remove(fn)
+
+
+@pytest.mark.parametrize('comp,ext', [(3, 'png')])
+@pytest.mark.parametrize('resx,resy', [(91, 81)])
+@pytest.mark.parametrize('dt', [ti.u16, ti.u32, ti.u64])
+@ti.host_arch_only
+def test_image_io_uint(resx, resy, comp, ext, dt):
+    shape = (resx, resy)
+    np_type = ti.to_numpy_type(dt)
+    # When saving to disk, pixel data will be truncated into 8 bits.
+    # Be careful here if you want lossless saving.
+    np_max = np.iinfo(np_type).max // 256
+    pixel = np.random.randint(256, size=(*shape, comp), dtype=np_type) * np_max
+    pixel_t = ti.Vector(comp, dt, shape)
+    pixel_t.from_numpy(pixel)
+    fn = make_temp(suffix='.' + ext)
+    ti.imwrite(pixel_t, fn)
+    pixel_r = ti.imread(fn).astype(np_type) * np_max
+    assert (pixel_r == pixel).all()
+    os.remove(fn)