faisalhr1997/codeformer

1# Modified from https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py  # noqa: E5012from abc import ABCMeta, abstractmethod3 4 5class BaseStorageBackend(metaclass=ABCMeta):6    """Abstract class of storage backends.7 8    All backends need to implement two apis: ``get()`` and ``get_text()``.9    ``get()`` reads the file as a byte stream and ``get_text()`` reads the file10    as texts.11    """12 13    @abstractmethod14    def get(self, filepath):15        pass16 17    @abstractmethod18    def get_text(self, filepath):19        pass20 21 22class MemcachedBackend(BaseStorageBackend):23    """Memcached storage backend.24 25    Attributes:26        server_list_cfg (str): Config file for memcached server list.27        client_cfg (str): Config file for memcached client.28        sys_path (str | None): Additional path to be appended to `sys.path`.29            Default: None.30    """31 32    def __init__(self, server_list_cfg, client_cfg, sys_path=None):33        if sys_path is not None:34            import sys35            sys.path.append(sys_path)36        try:37            import mc38        except ImportError:39            raise ImportError('Please install memcached to enable MemcachedBackend.')40 41        self.server_list_cfg = server_list_cfg42        self.client_cfg = client_cfg43        self._client = mc.MemcachedClient.GetInstance(self.server_list_cfg, self.client_cfg)44        # mc.pyvector servers as a point which points to a memory cache45        self._mc_buffer = mc.pyvector()46 47    def get(self, filepath):48        filepath = str(filepath)49        import mc50        self._client.Get(filepath, self._mc_buffer)51        value_buf = mc.ConvertBuffer(self._mc_buffer)52        return value_buf53 54    def get_text(self, filepath):55        raise NotImplementedError56 57 58class HardDiskBackend(BaseStorageBackend):59    """Raw hard disks storage backend."""60 61    def get(self, filepath):62        filepath = str(filepath)63        with open(filepath, 'rb') as f:64            value_buf = f.read()65        return value_buf66 67    def get_text(self, filepath):68        filepath = str(filepath)69        with open(filepath, 'r') as f:70            value_buf = f.read()71        return value_buf72 73 74class LmdbBackend(BaseStorageBackend):75    """Lmdb storage backend.76 77    Args:78        db_paths (str | list[str]): Lmdb database paths.79        client_keys (str | list[str]): Lmdb client keys. Default: 'default'.80        readonly (bool, optional): Lmdb environment parameter. If True,81            disallow any write operations. Default: True.82        lock (bool, optional): Lmdb environment parameter. If False, when83            concurrent access occurs, do not lock the database. Default: False.84        readahead (bool, optional): Lmdb environment parameter. If False,85            disable the OS filesystem readahead mechanism, which may improve86            random read performance when a database is larger than RAM.87            Default: False.88 89    Attributes:90        db_paths (list): Lmdb database path.91        _client (list): A list of several lmdb envs.92    """93 94    def __init__(self, db_paths, client_keys='default', readonly=True, lock=False, readahead=False, **kwargs):95        try:96            import lmdb97        except ImportError:98            raise ImportError('Please install lmdb to enable LmdbBackend.')99 100        if isinstance(client_keys, str):101            client_keys = [client_keys]102 103        if isinstance(db_paths, list):104            self.db_paths = [str(v) for v in db_paths]105        elif isinstance(db_paths, str):106            self.db_paths = [str(db_paths)]107        assert len(client_keys) == len(self.db_paths), ('client_keys and db_paths should have the same length, '108                                                        f'but received {len(client_keys)} and {len(self.db_paths)}.')109 110        self._client = {}111        for client, path in zip(client_keys, self.db_paths):112            self._client[client] = lmdb.open(path, readonly=readonly, lock=lock, readahead=readahead, **kwargs)113 114    def get(self, filepath, client_key):115        """Get values according to the filepath from one lmdb named client_key.116 117        Args:118            filepath (str | obj:`Path`): Here, filepath is the lmdb key.119            client_key (str): Used for distinguishing differnet lmdb envs.120        """121        filepath = str(filepath)122        assert client_key in self._client, (f'client_key {client_key} is not ' 'in lmdb clients.')123        client = self._client[client_key]124        with client.begin(write=False) as txn:125            value_buf = txn.get(filepath.encode('ascii'))126        return value_buf127 128    def get_text(self, filepath):129        raise NotImplementedError130 131 132class FileClient(object):133    """A general file client to access files in different backend.134 135    The client loads a file or text in a specified backend from its path136    and return it as a binary file. it can also register other backend137    accessor with a given name and backend class.138 139    Attributes:140        backend (str): The storage backend type. Options are "disk",141            "memcached" and "lmdb".142        client (:obj:`BaseStorageBackend`): The backend object.143    """144 145    _backends = {146        'disk': HardDiskBackend,147        'memcached': MemcachedBackend,148        'lmdb': LmdbBackend,149    }150 151    def __init__(self, backend='disk', **kwargs):152        if backend not in self._backends:153            raise ValueError(f'Backend {backend} is not supported. Currently supported ones'154                             f' are {list(self._backends.keys())}')155        self.backend = backend156        self.client = self._backends[backend](**kwargs)157 158    def get(self, filepath, client_key='default'):159        # client_key is used only for lmdb, where different fileclients have160        # different lmdb environments.161        if self.backend == 'lmdb':162            return self.client.get(filepath, client_key)163        else:164            return self.client.get(filepath)165 166    def get_text(self, filepath):167        return self.client.get_text(filepath)168