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)