-
Notifications
You must be signed in to change notification settings - Fork 1
/
rw_utils.py
44 lines (34 loc) · 1.31 KB
/
rw_utils.py
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
41
42
43
44
import os
import pickle
import numpy as np
from typing import Union
def serialize_boolean_array(array: np.array) -> bytes:
"""
Adapted from https://stackoverflow.com/a/48486714
Takes a np.ndarray with boolean values and converts it to a space-efficient
binary representation.
"""
return np.packbits(array).tobytes()
def deserialize_boolean_array(serialized_array: bytes, shape: tuple) -> np.ndarray:
"""
Adapted from https://stackoverflow.com/a/48486714
Inverse of serialize_boolean_array.
"""
num_elements = np.prod(shape)
packed_bits = np.frombuffer(serialized_array, dtype='uint8')
result = np.unpackbits(packed_bits)[:num_elements]
result.shape = shape
return result
def pickle_obj(obj, filename, path: Union[str, bytes, os.PathLike]) -> None:
if not os.path.isdir(path):
os.makedirs(path)
if not filename.endswith('.pickle'):
filename = f'{filename}.pickle'
with open(os.path.join(path, filename), 'wb') as wf:
pickle.dump(obj, file=wf, protocol=pickle.DEFAULT_PROTOCOL)
def read_pickled(filename: str, path: Union[str, bytes, os.PathLike]) -> object:
if not filename.endswith('.pickle'):
filename = f'{filename}.pickle'
with open(os.path.join(path, filename), 'rb') as rf:
obj = pickle.load(rf)
return obj