Skip to content

GitLab

Commit a6857292 authored by Vladislav Sovrasov's avatar Vladislav Sovrasov Committed by Kai Chen
Browse files

Fix acces to a wrong key in checkpoint (#734)

parent 0d9dbcc4
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
mmdet/apis/inference.py
View file @ a6857292
...@@ -34,7 +34,7 @@ def init_detector(config, checkpoint=None, device='cuda:0'): ...@@ -34,7 +34,7 @@ def init_detector(config, checkpoint=None, device='cuda:0'):
if checkpoint is not None: if checkpoint is not None:
checkpoint = load_checkpoint(model, checkpoint) checkpoint = load_checkpoint(model, checkpoint)
if 'CLASSES' in checkpoint['meta']: if 'CLASSES' in checkpoint['meta']:
model.CLASSES = checkpoint['meta']['classes'] model.CLASSES = checkpoint['meta']['CLASSES']
else: else:
warnings.warn('Class names are not saved in the checkpoint\'s ' warnings.warn('Class names are not saved in the checkpoint\'s '
'meta data, use COCO classes by default.') 'meta data, use COCO classes by default.')
......
tools/train.py
View file @ a6857292
...@@ -79,7 +79,7 @@ def main(): ...@@ -79,7 +79,7 @@ def main():
cfg.checkpoint_config.meta = dict( cfg.checkpoint_config.meta = dict(
mmdet_version=__version__, mmdet_version=__version__,
config=cfg.text, config=cfg.text,
classes=train_dataset.CLASSES) CLASSES=train_dataset.CLASSES)
# add an attribute for visualization convenience # add an attribute for visualization convenience
model.CLASSES = train_dataset.CLASSES model.CLASSES = train_dataset.CLASSES
train_detector( train_detector(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment