[WB-3854]: Out of range error when class names not included as argument (wandb#1479)

KyleGoyette · web-flow · commit 4150edbeecec · 2020-11-13T16:01:03.000-08:00
* fix conf mat code for now class names, add test

* staceys neat union cardinality fix

* handle edge case where inds are weird
diff --git a/standalone_tests/tweets.py b/standalone_tests/tweets.py
@@ -53,5 +53,11 @@
     class_ind_map[class_name] = i
 y_pred_inds = [class_ind_map[class_name] for class_name in y_pred]
 y_true_inds = [class_ind_map[class_name] for class_name in y_test]
-
+# test workflow with classes
 wandb.log({'conf_mat': wandb.plot.confusion_matrix(y_pred_inds, y_true_inds, nb.classes_)})
+# test workflow without classes
+wandb.log({'conf_mat_noclass': wandb.plot.confusion_matrix(y_pred_inds, y_true_inds)})
+# test workflow with multiples of inds
+y_pred_mult = [y_pred_ind*5 for y_pred_ind in y_pred_inds]
+y_true_mult = [y_true_ind*5 for y_true_ind in y_true_inds]
+wandb.log({'conf_mat_noclass_mult': wandb.plot.confusion_matrix(y_pred_mult, y_true_mult)})
diff --git a/wandb/plot/confusion_matrix.py b/wandb/plot/confusion_matrix.py
@@ -30,19 +30,26 @@ def confusion_matrix(preds=None, y_true=None, class_names=None):
     ), "Number of predictions and label indices must match"
     if class_names is not None:
         n_classes = len(class_names)
+        class_inds = set(preds).union(set(y_true))
         assert max(preds) <= len(
             class_names
         ), "Higher predicted index than number of classes"
         assert max(y_true) <= len(
             class_names
         ), "Higher label class index than number of classes"
     else:
-        n_classes = max(max(preds), max(y_true))
+        class_inds = set(preds).union(set(y_true))
+        n_classes = len(class_inds)
         class_names = ["Class_{}".format(i) for i in range(1, n_classes + 1)]
 
+    # get mapping of inds to class index in case user has weird prediction indices
+    class_mapping = {}
+    for i, val in enumerate(sorted(list(class_inds))):
+        class_mapping[val] = i
+
     counts = np.zeros((n_classes, n_classes))
     for i in range(len(preds)):
-        counts[y_true[i], preds[i]] += 1
+        counts[class_mapping[y_true[i]], class_mapping[preds[i]]] += 1
 
     data = []
     for i in range(n_classes):
