38 lines
973 B
Python
38 lines
973 B
Python
import os
|
|
import pickle
|
|
|
|
|
|
def get_dataset_path(lmdb_base_dir: str,
|
|
dataset_name: str) -> str:
|
|
"""
|
|
Get the path to the dataset in the lmdb directory
|
|
Args:
|
|
lmdb_base_dir: base directory of the lmdb dataset
|
|
dataset_name: name of the dataset
|
|
|
|
Returns:
|
|
path to the dataset
|
|
"""
|
|
|
|
dataset_path = os.path.join(lmdb_base_dir, dataset_name)
|
|
if not os.path.exists(dataset_path):
|
|
raise FileNotFoundError(f"Dataset {dataset_name} not found in {lmdb_base_dir}")
|
|
return dataset_path
|
|
|
|
|
|
def load_key_stats(lmdb_dir: str) -> dict:
|
|
"""
|
|
Load key statistics from the LMDB database.
|
|
:param lmdb_dir: Directory of the LMDB database.
|
|
:return: Dictionary with statistics.
|
|
"""
|
|
|
|
key_stats_path = f"{lmdb_dir}/key_stats.pickle"
|
|
|
|
if not os.path.exists(key_stats_path):
|
|
return dict()
|
|
|
|
with open(key_stats_path, "rb") as f:
|
|
key_stats = pickle.load(f)
|
|
return key_stats
|