import os def load_class_map(filename, root=''): class_map_path = filename if not os.path.exists(class_map_path): class_map_path = os.path.join(root, filename) assert os.path.exists(class_map_path), 'Cannot locate specified class map file (%s)' % filename class_map_ext = os.path.splitext(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)} else: assert False, 'Unsupported class map extension' return class_to_idx