def transform(dataset: IDataset, method: Union[str, Type[Transform]], *, env: Optional[Environment] = None, **kwargs) -> IDataset: """ Applies some function to dataset items. Results are computed lazily, if the transform supports this. Args: dataset - The dataset to be transformed method - The transformation to be applied to the dataset. If a string is passed, it is treated as a plugin name, which is searched for in the environment set by the 'env' argument env - A plugin collection. If not set, the built-in plugins are used **kwargs - Parameters for the transformation Returns: a wrapper around the input dataset """ if isinstance(method, str): if env is None: env = Environment() method = env.transforms[method] if not (inspect.isclass(method) and issubclass(method, Transform)): raise TypeError("Unexpected 'method' argument type: %s" % \ type(method)) produced = method(dataset, **kwargs) return Dataset(source=produced, env=env)
def test_can_extract_coco(self): tfds_example = { 'image': encode_image(np.ones((20, 10)), '.png'), 'image/filename': 'test.png', 'image/id': 123, 'objects': { 'bbox': [[0.1, 0.2, 0.3, 0.4]], 'label': [5], 'is_crowd': [True], } } with mock_tfds_data(example=tfds_example): tfds_info = tfds.builder('coco/2014').info expected_dataset = Dataset.from_iterable( [ DatasetItem( id='test', subset='train', image=np.ones((20, 10)), annotations=[ Bbox(2, 2, 2, 4, label=5, attributes={'is_crowd': True}), ], attributes={'id': 123}, ), ], categories=tfds_info.features['objects'].feature['label'].names ) extractor = make_tfds_extractor('coco/2014') actual_dataset = Dataset(extractor) compare_datasets(self, expected_dataset, actual_dataset, require_images=True)