add meta

Author: xvjiarui

tools/test.py

@@ -122,8 +122,16 @@ def main():
     if fp16_cfg is not None:
         wrap_fp16_model(model)
     checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
-    model.CLASSES = checkpoint['meta']['CLASSES']
-    model.PALETTE = checkpoint['meta']['PALETTE']
+    if 'CLASSES' in checkpoint.get('meta', {}):
+        model.CLASSES = checkpoint['meta']['CLASSES']
+    else:
+        print('"CLASSES" not found in meta, use dataset.CLASSES instead')
+        model.CLASSES = dataset.CLASSES
+    if 'PALETTE' in checkpoint.get('meta', {}):
+        model.PALETTE = checkpoint['meta']['PALETTE']
+    else:
+        print('"PALETTE" not found in meta, use dataset.PALETTE instead')
+        model.PALETTE = dataset.PALETTE
 
     efficient_test = False
     if args.eval_options is not None:

tools/train.py

@@ -149,6 +149,8 @@ def main():
             PALETTE=datasets[0].PALETTE)
     # add an attribute for visualization convenience
     model.CLASSES = datasets[0].CLASSES
+    # passing checkpoint meta for saving best checkpoint
+    meta.update(cfg.checkpoint_config.meta)
     train_segmentor(
         model,
         datasets,