def test_segmentation_select_masks():
data = [make_batch_select_masks(batch_size, vertex_ndim=2, npolygons_range=(1, 5),
nvertices_range=(3, 10)) for _ in range(data_size)]
pipe = segmentation_select_masks_pipeline(data)
compare_eager_with_pipeline(pipe, eager.segmentation.select_masks, eager_source=PipelineInput(
segmentation_select_masks_input_pipeline, data))
def get_data_source(*args, **kwargs):
return lambda: make_batch_select_masks(*args, **kwargs)