forked from CSAILVision/semantic-segmentation-pytorch
-
Notifications
You must be signed in to change notification settings - Fork 2
/
io_utils.py
49 lines (38 loc) · 1.22 KB
/
io_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
45
46
47
48
49
#!/usr/bin/env python3
import pickle
import gzip
import os
def _get_open_lambda(filename, is_save):
'''
with file extension pklz, we compress the pickle using gzip.
'''
extension = os.path.splitext(filename)[1]
compress = extension == '.pklz'
if compress:
open_mode = 'w' if is_save else 'r'
# note: I tried using bz2 instead of gzip. It compresses better
# but it is slower.
def open_lambda(x):
return gzip.GzipFile(x, open_mode)
else:
open_mode = 'wb' if is_save else 'rb'
def open_lambda(x):
return open(x, open_mode)
return open_lambda
def load_pickle(filename):
'''
load a pickle file
'''
open_lambda = _get_open_lambda(filename, is_save=False)
with open_lambda(filename) as file:
data = pickle.load(file)
return data
def save_pickle(filename, data):
'''
save a pickle file.
Note that pickle is incompatible across Python versions 2.7.x and 3.x
with file extension pklz, we compress the pickle using gzip
'''
open_lambda = _get_open_lambda(filename, is_save=True)
with open_lambda(filename) as file:
pickle.dump(data, file, protocol=pickle.HIGHEST_PROTOCOL)