update to support pickle based dictionaries

pull/1476/head
Alex Fafard 2 years ago
parent 1199c5a1a4
commit 7327792f39

@ -1,5 +1,5 @@
import os import os
import pickle
def load_class_map(map_or_filename, root=''): def load_class_map(map_or_filename, root=''):
if isinstance(map_or_filename, dict): if isinstance(map_or_filename, dict):
@ -13,6 +13,9 @@ def load_class_map(map_or_filename, root=''):
if class_map_ext == '.txt': if class_map_ext == '.txt':
with open(class_map_path) as f: with open(class_map_path) as f:
class_to_idx = {v.strip(): k for k, v in enumerate(f)} class_to_idx = {v.strip(): k for k, v in enumerate(f)}
elif class_map_ext == '.pkl':
with open(class_map_path,'rb') as f:
class_to_idx = pickle.load(f)
else: else:
assert False, f'Unsupported class map file extension ({class_map_ext}).' assert False, f'Unsupported class map file extension ({class_map_ext}).'
return class_to_idx return class_to_idx

Loading…
Cancel
Save