mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

Merge pull request #1476 from Afafard/class_dict

Browse Source 
update  to support pickle based dictionaries
This commit is contained in:
Ross Wightman 2022-09-27 13:26:34 -07:00 committed by GitHub
commit d4ea5c7d7d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
timm/data/parsers/class_map.py
View File

@ -1,5 +1,5 @@
import os
import pickle
def load_class_map(map_or_filename, root=''):
if isinstance(map_or_filename, dict):
@ -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