Skip to content

Commit

Permalink
Merge pull request #1476 from Afafard/class_dict
Browse files Browse the repository at this point in the history
update  to support pickle based dictionaries
  • Loading branch information
rwightman authored Sep 27, 2022
2 parents 1199c5a + 7327792 commit d4ea5c7
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion timm/data/parsers/class_map.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import os

import pickle

def load_class_map(map_or_filename, root=''):
if isinstance(map_or_filename, dict):
Expand All @@ -13,6 +13,9 @@ def load_class_map(map_or_filename, root=''):
if class_map_ext == '.txt':
with open(class_map_path) as 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:
assert False, f'Unsupported class map file extension ({class_map_ext}).'
return class_to_idx
Expand Down

0 comments on commit d4ea5c7

Please sign in to comment.