| 1234567891011121314151617181920212223 |
- import os
- import pickle
- def load_class_map(map_or_filename, root=''):
- if isinstance(map_or_filename, dict):
- assert dict, 'class_map dict must be non-empty'
- return map_or_filename
- class_map_path = map_or_filename
- if not os.path.exists(class_map_path):
- class_map_path = os.path.join(root, class_map_path)
- assert os.path.exists(class_map_path), 'Cannot locate specified class map file (%s)' % map_or_filename
- class_map_ext = os.path.splitext(map_or_filename)[-1].lower()
- 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
|
