Read a LMDB database and produce (k,v) raw bytes pairs. The raw bytes are usually not what you're interested in. You might want to use :class:`LMDBDataDecoder` or apply a mapper function after :class:`LMDBData`.
| 61 | |
| 62 | |
| 63 | class LMDBData(RNGDataFlow): |
| 64 | """ |
| 65 | Read a LMDB database and produce (k,v) raw bytes pairs. |
| 66 | The raw bytes are usually not what you're interested in. |
| 67 | You might want to use |
| 68 | :class:`LMDBDataDecoder` or apply a |
| 69 | mapper function after :class:`LMDBData`. |
| 70 | """ |
| 71 | def __init__(self, lmdb_path, shuffle=True, keys=None): |
| 72 | """ |
| 73 | Args: |
| 74 | lmdb_path (str): a directory or a file. |
| 75 | shuffle (bool): shuffle the keys or not. |
| 76 | keys (list[str] or str): list of str as the keys, used only when shuffle is True. |
| 77 | It can also be a format string e.g. ``{:0>8d}`` which will be |
| 78 | formatted with the indices from 0 to *total_size - 1*. |
| 79 | |
| 80 | If not given, it will then look in the database for ``__keys__`` which |
| 81 | :func:`LMDBSerializer.save` used to store the list of keys. |
| 82 | If still not found, it will iterate over the database to find |
| 83 | all the keys. |
| 84 | """ |
| 85 | self._lmdb_path = lmdb_path |
| 86 | self._shuffle = shuffle |
| 87 | |
| 88 | self._open_lmdb() |
| 89 | self._size = self._txn.stat()['entries'] |
| 90 | self._set_keys(keys) |
| 91 | logger.info("Found {} entries in {}".format(self._size, self._lmdb_path)) |
| 92 | |
| 93 | # Clean them up after finding the list of keys, since we don't want to fork them |
| 94 | self._close_lmdb() |
| 95 | |
| 96 | def _set_keys(self, keys=None): |
| 97 | def find_keys(txn, size): |
| 98 | logger.warn("Traversing the database to find keys is slow. Your should specify the keys.") |
| 99 | keys = [] |
| 100 | with timed_operation("Loading LMDB keys ...", log_start=True), \ |
| 101 | get_tqdm(total=size) as pbar: |
| 102 | for k in self._txn.cursor(): |
| 103 | assert k[0] != b'__keys__' |
| 104 | keys.append(k[0]) |
| 105 | pbar.update() |
| 106 | return keys |
| 107 | |
| 108 | self.keys = self._txn.get(b'__keys__') |
| 109 | if self.keys is not None: |
| 110 | self.keys = loads(self.keys) |
| 111 | self._size -= 1 # delete this item |
| 112 | |
| 113 | if self._shuffle: # keys are necessary when shuffle is True |
| 114 | if keys is None: |
| 115 | if self.keys is None: |
| 116 | self.keys = find_keys(self._txn, self._size) |
| 117 | else: |
| 118 | # check if key-format like '{:0>8d}' was given |
| 119 | if isinstance(keys, six.string_types): |
| 120 | self.keys = map(lambda x: keys.format(x), list(np.arange(self._size))) |