code: def check_select_masks(
      batch_size,
      npolygons_range=(1, 10),
      nvertices_range=(3, 40),
      vertex_ndim=2,
      vertex_dtype=np.float32,
      reindex_masks=False,
  ):
      def get_data_source(*args, **kwargs):
          return lambda: make_batch_select_masks(*args, **kwargs)
  
      pipe = dali.pipeline.Pipeline(
          batch_size=batch_size, num_threads=4, device_id=0, seed=1234
      )
      with pipe:
          polygons, vertices, mask_ids = fn.external_source(
              source=get_data_source(
                  batch_size,
                  npolygons_range=npolygons_range,
                  nvertices_range=nvertices_range,
                  vertex_ndim=vertex_ndim,
                  vertex_dtype=vertex_dtype,
              ),
              num_outputs=3,
              device="cpu",
          )
          out_polygons, out_vertices = fn.segmentation.select_masks(
              mask_ids, polygons, vertices, reindex_masks=reindex_masks
          )
      pipe.set_outputs(
          polygons, vertices, mask_ids, out_polygons, out_vertices
      )
      pipe.build()
      for iter in range(3):
          outputs = pipe.run()
          for idx in range(batch_size):
              in_polygons = outputs[0].at(idx)
              in_vertices = outputs[1].at(idx)
              mask_ids = outputs[2].at(idx)
              out_polygons = outputs[3].at(idx)
              out_vertices = outputs[4].at(idx)
  
              in_polygons_dict = {}
              for k in range(in_polygons.shape[0]):
                  mask_id = in_polygons[k, 0]
                  in_polygons_dict[mask_id] = (
                      in_polygons[k, 1],
                      in_polygons[k, 2],
                  )
  
              if reindex_masks:
                  index_map = {}
                  for idx in range(len(mask_ids)):
                      index_map[mask_ids[idx]] = idx
  
              vertex_count = 0
              for m in range(len(mask_ids)):
                  mask_id = mask_ids[m]
                  in_vertex_start, in_vertex_end = in_polygons_dict[
                      mask_id
                  ]
                  in_nvertices = in_vertex_end - in_vertex_start
  
                  expected_out_mask_id = (
                      index_map[mask_id] if reindex_masks else mask_id
                  )
                  (
                      out_mask_id,
                      out_vertex_start,
                      out_vertex_end,
                  ) = out_polygons[m]
                  assert out_mask_id == expected_out_mask_id
                  assert out_vertex_start == vertex_count
                  assert out_vertex_end == (vertex_count + in_nvertices)
                  vertex_count = vertex_count + in_nvertices
                  expected_out_vertex = in_vertices[
                      in_vertex_start:in_vertex_end
                  ]
                  out_vertex = out_vertices[
                      out_vertex_start:out_vertex_end
                  ]
                  assert (expected_out_vertex == out_vertex).all()
  
file_path: DALI/dali/test/python/operator_2/test_segmentation_select_masks.py
imports: from nose_utils import assert_raises
  import numpy
  import nvidia.dali
  import nvidia.dali.fn
  import random
  from segmentation_test_utils import make_batch_select_masks
label_desc: Use image node_id 1 for calling a global function with example usage: check_select_masks(batch_size, npolygons_range, nvertices_range, vertex_ndim, vertex_dtype, reindex_masks) and returns: function