Add ConfusionMatrix normalize=True flag (ultralytics#3586)

fhkiel-mlaip · Jun 11, 2021 · f3413e9 · f3413e9
1 parent b3d7fc2
commit f3413e9
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/utils/metrics.py b/utils/metrics.py
@@ -158,11 +158,12 @@ def process_batch(self, detections, labels):
     def matrix(self):
         return self.matrix
 
-    def plot(self, save_dir='', names=()):
+    def plot(self, normalize=True, save_dir='', names=()):
         try:
             import seaborn as sn
-
-            array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6)  # normalize
+
+            if normalize:
+                array = self.matrix / (self.matrix.sum(0).reshape(1, self.nc + 1) + 1E-6)  # normalize columns
             array[array < 0.005] = np.nan  # don't annotate (would appear as 0.00)
 
             fig = plt.figure(figsize=(12, 9), tight_layout=True)