[common] (ImagePipelineEvaluation) add labels to confusion matrix

common/research_common/image_pipeline_evaluation/image_pipeline_evaluation_reporter.py

@@ -2,20 +2,20 @@ from collections import Counter
 from dataclasses import dataclass
 from os.path import relpath
 from pathlib import Path
-from typing import Iterable, List, Any, Dict, Tuple
+from typing import Any, Dict, Iterable, List, Tuple
 
 import numpy as np
 from pandas import DataFrame
 
 from polystar.common.image_pipeline.image_pipeline import ImagePipeline
-from polystar.common.utils.dataframe import format_df_rows, format_df_row, format_df_column
+from polystar.common.utils.dataframe import format_df_column, format_df_row, format_df_rows
 from polystar.common.utils.markdown import MarkdownFile
 from polystar.common.utils.time import create_time_id
-from research_common.constants import EVALUATION_DIR, DSET_DIR
+from research_common.constants import DSET_DIR, EVALUATION_DIR
 from research_common.dataset.roco_dataset import ROCODataset
 from research_common.image_pipeline_evaluation.image_pipeline_evaluator import (
-    ImagePipelineEvaluator,
     ClassificationResults,
+    ImagePipelineEvaluator,
     SetClassificationResults,
 )
 
@@ -104,7 +104,7 @@ class ImagePipelineEvaluationReporter:
         format_df_row(df, "support", int)
         mf.table(df)
         mf.title("Confusion Matrix:", level=4)
-        mf.table(DataFrame(results.confusion_matrix))
+        mf.table(DataFrame(results.confusion_matrix, index=results.unique_labels, columns=results.unique_labels))
         mf.title("10 Mistakes examples", level=4)
         mistakes_idx = np.random.choice(results.mistakes, min(len(results.mistakes), 10), replace=False)
         relative_paths = [

common/research_common/image_pipeline_evaluation/image_pipeline_evaluator.py

 import logging
 from dataclasses import dataclass
 from time import time
-from typing import List, Dict, Any, Iterable, Sequence
+from typing import Any, Dict, Iterable, List, Sequence
 
 import numpy as np
 from sklearn.metrics import classification_report, confusion_matrix
 
+from memoized_property import memoized_property
 from polystar.common.image_pipeline.image_pipeline import ImagePipeline
 from polystar.common.models.image import Image
 from research_common.dataset.directory_roco_dataset import DirectoryROCODataset
@@ -30,6 +31,10 @@ class SetClassificationResults:
     def mistakes(self) -> Sequence[int]:
         return np.where(self.labels != self.predictions)[0]
 
+    @memoized_property
+    def unique_labels(self) -> List[Any]:
+        return sorted(set(self.labels) | set(self.predictions))
+
 
 @dataclass
 class ClassificationResults: