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.field(comp, dt, shape)
pixel_t.from_numpy(pixel)
fn = make_temp_file(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)
def create_field(dim, dtype, shape=None, initial=None, **kwargs):
if dim is None:
return dtype(shape, **kwargs)
if dim == 0:
dim = ()
if not isinstance(dim, (list, tuple)):
dim = [dim]
if isinstance(dim, list):
dim = tuple(dim)
if len(dim) == 0:
ret = ti.field(dtype, shape, **kwargs)
elif len(dim) == 1:
ret = ti.Vector.field(dim[0], dtype, shape, **kwargs)
elif len(dim) == 2:
ret = ti.Matrix.field(dim[0], dim[1], dtype, shape, **kwargs)
else:
raise TypeError(f'Expect int or tuple for dim, got: {dim}')
if initial is not None:
if callable(initial):
ti.materialize_callback(initial)
else:
initial = np.array(initial, dtype=ti.to_numpy_type(dtype))
@ti.materialize_callback
def init_field():
ret.from_numpy(initial)
return ret
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.field(comp, dt, shape)
pixel_t.from_numpy(pixel)
fn = make_temp_file(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)
def test_image_io(resx, resy, comp, ext, is_field, dt):
if comp != 1:
shape = (resx, resy, comp)
else:
shape = (resx, resy)
if is_field:
pixel_t = ti.field(dt, shape)
pixel = np.random.randint(256, size=shape, dtype=ti.to_numpy_type(dt))
if is_field:
pixel_t.from_numpy(pixel)
fn = make_temp_file(suffix='.' + ext)
if is_field:
ti.imwrite(pixel_t, fn)
else:
ti.imwrite(pixel, 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)
def func():
ck = ti.to_numpy_type(vartype)(0)
for i in range(n):
x[i] = ti.atomic_add(c[None], step)
y[i] = ti.atomic_add(ck, step)
def to_numpy_type(dtype):
if dtype is int:
dtype = ti.get_runtime().default_ip
elif dtype is float:
dtype = ti.get_runtime().default_fp
return ti.to_numpy_type(dtype)
