ImageNet names, classify inference, resume fixes (#712)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com>
This commit is contained in:
@ -226,10 +226,13 @@ class AutoBackend(nn.Module):
|
||||
f"https://docs.ultralytics.com/reference/nn/")
|
||||
|
||||
# class names
|
||||
if 'names' not in locals():
|
||||
names = yaml_load(data)['names'] if data else {i: f'class{i}' for i in range(999)}
|
||||
if names[0] == 'n01440764' and len(names) == 1000: # ImageNet
|
||||
names = yaml_load(ROOT / 'yolo/data/datasets/ImageNet.yaml')['names'] # human-readable names
|
||||
if 'names' not in locals(): # names missing
|
||||
names = yaml_load(data)['names'] if data else {i: f'class{i}' for i in range(999)} # assign default
|
||||
elif isinstance(names, list): # names is a list
|
||||
names = dict(enumerate(names)) # convert to dict
|
||||
if isinstance(names[0], str) and names[0].startswith('n0'): # imagenet class codes, i.e. 'n01440764'
|
||||
map = yaml_load(ROOT / 'yolo/data/datasets/ImageNet.yaml')['map'] # human-readable names
|
||||
names = {k: map[v] for k, v in names.items()}
|
||||
|
||||
self.__dict__.update(locals()) # assign all variables to self
|
||||
|
||||
|
Reference in New Issue
Block a user