Function test_video_decoder_frame_indices

dali/test/python/decoder/test_video.py:519–568

(device, batch_size, filenames, frames, pad_mode)

    ],
)
def test_video_decoder_frame_indices(device, batch_size, filenames, frames, pad_mode):
    skip_if_m60()
    num_iters = 1
    batch = []
    for i in range(batch_size):
        with open(filenames[i % len(filenames)], "rb") as f:
            batch.append(np.frombuffer(f.read(), dtype=np.uint8))

    def get_batch():
        return batch

    fill_value = [118, 185, 0]

    @pipeline_def
    def test_pipeline():
        encoded = fn.external_source(source=get_batch, device="cpu")
        reference = fn.experimental.decoders.video(encoded, device=device)
        if frames is None:
            frames_shape = fn.random.uniform(range=(1, 10), shape=(1,), dtype=types.INT32)
            frames_arg = fn.random.uniform(range=(0, 10), shape=frames_shape, dtype=types.INT32)
        else:
            frames_arg = frames
        decoded = fn.experimental.decoders.video(
            encoded, device=device, frames=frames_arg, pad_mode=pad_mode, fill_value=fill_value
        )
        return (reference, decoded, frames_arg)

    pipe = test_pipeline(batch_size=batch_size, num_threads=3, device_id=0)
    for _ in range(num_iters):
        out0, out1, frames_arg = (o.as_cpu() for o in pipe.run())
        # Verify each sample in the batch
        for i in range(batch_size):
            full_video = out0.at(i)
            F, H, W, C = full_video.shape
            out = out1.at(i)
            frames = frames_arg.at(i)
            # Get frames at specified indices for reference
            if pad_mode == "none":
                for f in frames:
                    assert f < full_video.shape[0]
                out_reference = full_video[frames]
            else:
                padding = np.full((H, W, C), fill_value, dtype=np.uint8)
                out_reference = np.zeros((len(frames), H, W, C), dtype=np.uint8)
                for f in range(len(frames)):
                    if frames[f] < F:
                        out_reference[f] = full_video[frames[f]]
                    else:
                        out_reference[f] = padding
            compare_videos(out_reference, out)


@params(