[Feature] Add confusion matrix to allowed metrics in BaseDataset

mmaction/datasets/base.py

@@ -11,8 +11,10 @@
 from mmcv.utils import print_log
 from torch.utils.data import Dataset
 
-from ..core import (mean_average_precision, mean_class_accuracy,
-                    mmit_mean_average_precision, top_k_accuracy)
+from ..core import (confusion_matrix, mean_average_precision,
+                    mean_class_accuracy, mmit_mean_average_precision,
+                    top_k_accuracy)
+from ..utils import visualize_confusion_matrix
 from .pipelines import Compose
 
 
@@ -179,7 +181,7 @@ def evaluate(self,
         metrics = metrics if isinstance(metrics, (list, tuple)) else [metrics]
         allowed_metrics = [
             'top_k_accuracy', 'mean_class_accuracy', 'mean_average_precision',
-            'mmit_mean_average_precision'
+            'mmit_mean_average_precision', 'confusion_matrix'
         ]
 
         for metric in metrics:
@@ -240,6 +242,13 @@ def evaluate(self,
                 print_log(log_msg, logger=logger)
                 continue
 
+            if metric == 'confusion_matrix':
+                y_pred = [np.argmax(result) for result in results]
+                cm = confusion_matrix(y_pred, gt_labels, 'true')
+                print_log(cm, logger=logger)
+                visualize_confusion_matrix(cm)
+                continue
+
         return eval_results
 
     @staticmethod

mmaction/utils/__init__.py

@@ -2,13 +2,14 @@
 from .collect_env import collect_env
 from .gradcam_utils import GradCAM
 from .logger import get_root_logger
-from .misc import get_random_string, get_shm_dir, get_thread_id
+from .misc import (get_random_string, get_shm_dir, get_thread_id,
+                   visualize_confusion_matrix)
 from .module_hooks import register_module_hooks
 from .precise_bn import PreciseBNHook
 from .setup_env import setup_multi_processes
 
 __all__ = [
     'get_root_logger', 'collect_env', 'get_random_string', 'get_thread_id',
     'get_shm_dir', 'GradCAM', 'PreciseBNHook', 'register_module_hooks',
-    'setup_multi_processes'
+    'setup_multi_processes', 'visualize_confusion_matrix'
 ]

mmaction/utils/misc.py

@@ -25,3 +25,17 @@ def get_thread_id():
 def get_shm_dir():
     """Get shm dir for temporary usage."""
     return '/dev/shm'
+
+
+def visualize_confusion_matrix(confusion_matrix):
+    """Visualize a confusion matrix.
+
+    Args:
+        confusion_matrix (np.array): the confusion matrix
+    """
+    import matplotlib.pyplot as plt
+    import seaborn as sns
+    plt.figure(figsize=(15, 10))
+    sns.set(font_scale=1.5)
+    sns.heatmap(confusion_matrix, annot=True, square=True, cbar=True)
+    plt.show()