display_matrix text color fix

MiXaiLL76 · Aug 29, 2024 · 77b8631 · 77b8631
1 parent 697204d
commit 77b8631
Showing 1 changed file with 12 additions and 20 deletions.
diff --git a/faster_coco_eval/extra/draw.py b/faster_coco_eval/extra/draw.py
@@ -306,33 +306,26 @@ def display_matrix(
     else:
         hovertemplate += "Count: %{z:.0f}<extra></extra>"
 
+    annotations = []
+    for j, row in enumerate(conf_matrix):
+        annotations.append([])
+        for value in row:
+            text_value = "{:.0f}".format(value)
+            if normalize:
+                text_value += "%"
+
+            annotations[j].append(text_value)
+
     heatmap = go.Heatmap(
         z=conf_matrix,
         x=_labels,
         y=_labels[:-2],
+        text=annotations,
         colorscale="Blues",
         hovertemplate=hovertemplate,
+        texttemplate="%{text}",
     )
 
-    annotations = []
-    for j, row in enumerate(conf_matrix):
-        for i, value in enumerate(row):
-            text_value = "{:.0f}".format(value)
-            if normalize:
-                text_value += "%"
-
-            annotations.append(
-                {
-                    "x": _labels[i],
-                    "y": _labels[j],
-                    "font": {"color": "white"},
-                    "text": text_value,
-                    "xref": "x1",
-                    "yref": "y1",
-                    "showarrow": False,
-                }
-            )
-
     title = "Confusion Matrix"
     if normalize:
         title = "Normalized " + title
@@ -341,7 +334,6 @@ def display_matrix(
         "title": title,
         "xaxis": {"title": "Predicted value"},
         "yaxis": {"title": "Real value"},
-        "annotations": annotations,
     }
 
     fig = go.Figure(data=[heatmap], layout=layout)