def METHOD_NAME(image_set, return_masks):
    prepare = ConvertCocoPolysToMask(return_masks)
    transforms = make_coco_transforms(image_set)
    def map(data):
        image, target = data
        numpy_image = image.numpy()
        pil_image = im.fromarray(numpy_image)
        coco_target = {'image_id': target.image_id, 'annotations': target.coco()}
        image, target = prepare(pil_image, coco_target)
        transformed_img, transformed_target = transforms(image, target)
        return transformed_img, transformed_target
    return map