def _run_test_stack(num_inputs, layout, ndim, axis, axis_name):
num_iter=3
batch_size=4
if ndim is None:
ndim = len(layout)
ref_axis = axis if axis is not None else 0
if axis_name:
ref_layout = layout[:axis] + axis_name + layout[axis:] if layout else axis_name
else:
ref_layout = ""
pipe = dali.pipeline.Pipeline(batch_size=batch_size, num_threads = 3, device_id = 0)
with pipe:
inputs = fn.external_source(input_generator(num_inputs, batch_size, ndim), num_outputs=num_inputs, layout=layout)
out_cpu = fn.stack(*inputs, axis=axis, axis_name=axis_name)
out_gpu = fn.stack(*(x.gpu() for x in inputs), axis=axis, axis_name=axis_name)
pipe.set_outputs(out_cpu, out_gpu, *inputs);
pipe.build()
for _ in range(num_iter):
o_cpu, o_gpu, *inputs = pipe.run()
ref = ref_stack(inputs, ref_axis)
check_batch(o_cpu, ref, batch_size, eps=0, expected_layout=ref_layout)
check_batch(o_gpu, ref, batch_size, eps=0, expected_layout=ref_layout)
def create_pipline():
jpegs = fn.external_source(source=callback,
num_outputs=sequence_lenght * 2,
parallel=parallel,
batch=False)
images = fn.decoders.image(jpegs, device="cpu")
sequence = fn.stack(*images)
sequence = fn.reshape(sequence, layout="DHWC")
return sequence
def test_stack_cpu():
pipe = Pipeline(batch_size=batch_size, num_threads=3, device_id=None)
data = fn.external_source(source=get_data, layout="HWC")
data2 = fn.external_source(source=get_data, layout="HWC")
data3 = fn.external_source(source=get_data, layout="HWC")
pixel_pos = fn.stack(data, data2, data3)
pipe.set_outputs(pixel_pos)
pipe.build()
for _ in range(3):
pipe.run()
def check_random_mask_pixel(ndim=2, batch_size=3,
min_extent=20, max_extent=50):
pipe = dali.pipeline.Pipeline(batch_size=batch_size, num_threads=4, device_id=0, seed=1234)
with pipe:
# Input mask
in_shape_dims = [fn.cast(fn.random.uniform(range=(min_extent, max_extent + 1)),
dtype=types.INT32) for _ in range(ndim)]
in_shape = fn.stack(*in_shape_dims)
in_mask = fn.cast(fn.random.uniform(range=(0, 2), shape=in_shape), dtype=types.INT32)
# > 0
fg_pixel1 = fn.segmentation.random_mask_pixel(in_mask, foreground=1)
# >= 0.99
fg_pixel2 = fn.segmentation.random_mask_pixel(in_mask, foreground=1, threshold=0.99)
# == 2
fg_pixel3 = fn.segmentation.random_mask_pixel(in_mask, foreground=1, value=2)
rnd_pixel = fn.segmentation.random_mask_pixel(in_mask, foreground=0)
coin_flip = fn.random.coin_flip(probability=0.7)
fg_biased = fn.segmentation.random_mask_pixel(in_mask, foreground=coin_flip)
# Demo purposes: Taking a random pixel and produce a valid anchor to feed slice
# We want to force the center adjustment, thus the large crop shape
crop_shape = in_shape - 2
anchor = fn.cast(fg_pixel1, dtype=types.INT32) - crop_shape // 2
anchor = math.min(math.max(0, anchor), in_shape - crop_shape)
out_mask = fn.slice(in_mask, anchor, crop_shape, axes=tuple(range(ndim)))
pipe.set_outputs(in_mask, fg_pixel1, fg_pixel2, fg_pixel3, rnd_pixel, coin_flip, fg_biased,
anchor, crop_shape, out_mask)
pipe.build()
for iter in range(3):
outputs = pipe.run()
for idx in range(batch_size):
in_mask = outputs[0].at(idx)
fg_pixel1 = outputs[1].at(idx).tolist()
fg_pixel2 = outputs[2].at(idx).tolist()
fg_pixel3 = outputs[3].at(idx).tolist()
rnd_pixel = outputs[4].at(idx).tolist()
coin_flip = outputs[5].at(idx).tolist()
fg_biased = outputs[6].at(idx).tolist()
anchor = outputs[7].at(idx).tolist()
crop_shape = outputs[8].at(idx).tolist()
out_mask = outputs[9].at(idx)
assert in_mask[tuple(fg_pixel1)] > 0
assert in_mask[tuple(fg_pixel2)] > 0.99
assert in_mask[tuple(fg_pixel3)] == 2
assert in_mask[tuple(fg_biased)] > 0 or not coin_flip
for d in range(ndim):
assert 0 <= anchor[d] and anchor[d] + crop_shape[d] <= in_mask.shape[d]
assert out_mask.shape == tuple(crop_shape)
