mirror of
https://github.com/huggingface/pytorch-image-models.git
synced 2025-06-03 15:01:08 +08:00
ba65dfe2c6
* support some torchvision datasets * improvements to TFDS wrapper for subsplit handling (fix #942), shuffle seed * add class-map support to train (fix #957)
20 lines
759 B
Python
20 lines
759 B
Python
import os
|
|
|
|
|
|
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)}
|
|
else:
|
|
assert False, f'Unsupported class map file extension ({class_map_ext}).'
|
|
return class_to_idx
|
|
|