Spaces:
Running
Running
File size: 1,213 Bytes
0fc5df9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import os
from random import choices, randint
from typing import cast, Dict, List, Optional, TypedDict
import h5py
datasets_dir: str = './datasets'
datasets_file: str = 'pregenerated_pokemon.h5'
h5_file: str = os.path.join(datasets_dir, datasets_file)
class Stats(TypedDict):
size_total: int
size_mb: float
size_counts: Dict[str, int]
def get_stats(h5_file: str = h5_file) -> Stats:
with h5py.File(h5_file, 'r') as datasets:
return {
"size_total": sum(list(datasets[energy].size.item() for energy in datasets.keys())),
"size_mb": round(os.path.getsize(h5_file) / 1024**2, 1),
"size_counts": {key: datasets[key].size.item() for key in datasets.keys()},
}
energy_types: List[str] = ['colorless', 'darkness', 'dragon', 'fairy', 'fighting',
'fire', 'grass', 'lightning', 'metal', 'psychic', 'water']
def get_image(energy: Optional[str] = None, row: Optional[int] = None) -> str:
if not energy:
energy = choices(energy_types)[0]
with h5py.File(h5_file, 'r') as datasets:
if not row:
row = randint(0, datasets[energy].size - 1)
return datasets[energy].asstr()[row][0]
|