Function get_input_arg_per_frame

dali/test/python/sequences_test_utils.py:327–345

(input_data: ArgData, param_cb, rng, check_broadcasting)

def get_input_arg_per_frame(input_data: ArgData, param_cb, rng, check_broadcasting):
    frame_idx = input_data.desc.expandable_prefix.find("F")
    assert frame_idx >= 0

    def arg_for_sample(sample_idx, batch_idx, sample):
        if check_broadcasting and rng.randint(1, 4) == 1:
            return np.array([param_cb(SampleDesc(rng, 0, sample_idx, batch_idx, sample))])
        num_frames = sample.shape[frame_idx]
        return np.array(
            [
                param_cb(SampleDesc(rng, frame_idx, sample_idx, batch_idx, sample))
                for frame_idx in range(num_frames)
            ]
        )

    return [
        [arg_for_sample(sample_idx, batch_idx, sample) for sample_idx, sample in enumerate(batch)]
        for batch_idx, batch in enumerate(input_data.data)
    ]


def compute_input_params_data(input_data: ArgData, rng, input_params: List[ArgCb]):