Spaces:
Paused
Paused
| # Copyright (c) OpenMMLab. All rights reserved. | |
| from io import BytesIO, StringIO | |
| from pathlib import Path | |
| from ..utils import is_list_of, is_str | |
| from .file_client import FileClient | |
| from .handlers import BaseFileHandler, JsonHandler, PickleHandler, YamlHandler | |
| file_handlers = { | |
| 'json': JsonHandler(), | |
| 'yaml': YamlHandler(), | |
| 'yml': YamlHandler(), | |
| 'pickle': PickleHandler(), | |
| 'pkl': PickleHandler() | |
| } | |
| def load(file, file_format=None, file_client_args=None, **kwargs): | |
| """Load data from json/yaml/pickle files. | |
| This method provides a unified api for loading data from serialized files. | |
| Note: | |
| In v1.3.16 and later, ``load`` supports loading data from serialized | |
| files those can be storaged in different backends. | |
| Args: | |
| file (str or :obj:`Path` or file-like object): Filename or a file-like | |
| object. | |
| file_format (str, optional): If not specified, the file format will be | |
| inferred from the file extension, otherwise use the specified one. | |
| Currently supported formats include "json", "yaml/yml" and | |
| "pickle/pkl". | |
| file_client_args (dict, optional): Arguments to instantiate a | |
| FileClient. See :class:`mmcv.fileio.FileClient` for details. | |
| Default: None. | |
| Examples: | |
| >>> load('/path/of/your/file') # file is storaged in disk | |
| >>> load('https://path/of/your/file') # file is storaged in Internet | |
| >>> load('s3://path/of/your/file') # file is storaged in petrel | |
| Returns: | |
| The content from the file. | |
| """ | |
| if isinstance(file, Path): | |
| file = str(file) | |
| if file_format is None and is_str(file): | |
| file_format = file.split('.')[-1] | |
| if file_format not in file_handlers: | |
| raise TypeError(f'Unsupported format: {file_format}') | |
| handler = file_handlers[file_format] | |
| if is_str(file): | |
| file_client = FileClient.infer_client(file_client_args, file) | |
| if handler.str_like: | |
| with StringIO(file_client.get_text(file)) as f: | |
| obj = handler.load_from_fileobj(f, **kwargs) | |
| else: | |
| with BytesIO(file_client.get(file)) as f: | |
| obj = handler.load_from_fileobj(f, **kwargs) | |
| elif hasattr(file, 'read'): | |
| obj = handler.load_from_fileobj(file, **kwargs) | |
| else: | |
| raise TypeError('"file" must be a filepath str or a file-object') | |
| return obj | |
| def dump(obj, file=None, file_format=None, file_client_args=None, **kwargs): | |
| """Dump data to json/yaml/pickle strings or files. | |
| This method provides a unified api for dumping data as strings or to files, | |
| and also supports custom arguments for each file format. | |
| Note: | |
| In v1.3.16 and later, ``dump`` supports dumping data as strings or to | |
| files which is saved to different backends. | |
| Args: | |
| obj (any): The python object to be dumped. | |
| file (str or :obj:`Path` or file-like object, optional): If not | |
| specified, then the object is dumped to a str, otherwise to a file | |
| specified by the filename or file-like object. | |
| file_format (str, optional): Same as :func:`load`. | |
| file_client_args (dict, optional): Arguments to instantiate a | |
| FileClient. See :class:`mmcv.fileio.FileClient` for details. | |
| Default: None. | |
| Examples: | |
| >>> dump('hello world', '/path/of/your/file') # disk | |
| >>> dump('hello world', 's3://path/of/your/file') # ceph or petrel | |
| Returns: | |
| bool: True for success, False otherwise. | |
| """ | |
| if isinstance(file, Path): | |
| file = str(file) | |
| if file_format is None: | |
| if is_str(file): | |
| file_format = file.split('.')[-1] | |
| elif file is None: | |
| raise ValueError( | |
| 'file_format must be specified since file is None') | |
| if file_format not in file_handlers: | |
| raise TypeError(f'Unsupported format: {file_format}') | |
| handler = file_handlers[file_format] | |
| if file is None: | |
| return handler.dump_to_str(obj, **kwargs) | |
| elif is_str(file): | |
| file_client = FileClient.infer_client(file_client_args, file) | |
| if handler.str_like: | |
| with StringIO() as f: | |
| handler.dump_to_fileobj(obj, f, **kwargs) | |
| file_client.put_text(f.getvalue(), file) | |
| else: | |
| with BytesIO() as f: | |
| handler.dump_to_fileobj(obj, f, **kwargs) | |
| file_client.put(f.getvalue(), file) | |
| elif hasattr(file, 'write'): | |
| handler.dump_to_fileobj(obj, file, **kwargs) | |
| else: | |
| raise TypeError('"file" must be a filename str or a file-object') | |
| def _register_handler(handler, file_formats): | |
| """Register a handler for some file extensions. | |
| Args: | |
| handler (:obj:`BaseFileHandler`): Handler to be registered. | |
| file_formats (str or list[str]): File formats to be handled by this | |
| handler. | |
| """ | |
| if not isinstance(handler, BaseFileHandler): | |
| raise TypeError( | |
| f'handler must be a child of BaseFileHandler, not {type(handler)}') | |
| if isinstance(file_formats, str): | |
| file_formats = [file_formats] | |
| if not is_list_of(file_formats, str): | |
| raise TypeError('file_formats must be a str or a list of str') | |
| for ext in file_formats: | |
| file_handlers[ext] = handler | |
| def register_handler(file_formats, **kwargs): | |
| def wrap(cls): | |
| _register_handler(cls(**kwargs), file_formats) | |
| return cls | |
| return wrap | |