NNT Visualize Prediction Metrics

Overall accuracy can hide a disaster. A model that nails 95% of classes 1–9 but can't tell a 3 from an 8 still reports a great number. NNT Visualize Prediction Metrics is the node that surfaces exactly that: feed it the metrics dict from your evaluation step, and it draws the confusion matrix and per-class accuracy bars that show you where your model is fooled. It's the "who actually is my model confused about" node.

In the pack's flow, this takes the DICT produced by the evaluation side of a trained model - the one containing the confusion matrix, per-class accuracy, and overall accuracy (or, for regression, MSE/MAE with true vs predicted values). It detects which kind of problem you ran and plots accordingly.

How it works

The input is metrics (a DICT) plus a plot_type and the usual image_width/image_height. The three plot types:

• confusion_matrix - a seaborn heatmap of predicted vs true class counts. The diagonal lights up where the model is right; the bright off-diagonal cells are your systematic confusions. For MNIST, that "is it a 4 or a 9" cell will glow.
• class_accuracy - a bar chart of accuracy per class. Instantly shows you the weak classes.
• combined - both, stacked. The default, and usually all you need.

It's built to handle both classification and regression metrics. If the dict has confusion_matrix / per_class_accuracy / accuracy, you get the classifier views; if it has mse / mae / true_values / predictions, it switches to a regression-style plot comparing true against predicted values. Outputs: an IMAGE (the rendered figure, ready for Preview Image) and a STRING summary.

The gotchas

The metrics dict has to contain the keys the node expects, in the shapes it expects. The pack's own evaluation node produces exactly that (classification: accuracy, per_class_accuracy, confusion_matrix, predictions; regression: mse, mae, predictions, true_values) - but if you're hand-building a metrics dict, mismatch the keys and the node quietly draws whatever is there or errors. Also, the node prints the received metric keys to the console on every run - that's debugging output left in, not a bug, though it does make your console chatty.

One more thing: the confusion matrix values are numpy-array-backed lists, and the heatmap annotates with integer counts. For a 10-class MNIST model that's perfect; for a 100-class CIFAR100 model the matrix becomes a 100×100 blur you can't read. If you're training something with many classes, class_accuracy (or a small subset) is the more useful view.

Install

Pack-level install:

cd ComfyUI/custom_nodes
git clone https://github.com/inventorado/ComfyUI_NNT.git
cd ComfyUI_NNT
pip install -r requirements.txt

Restart ComfyUI (or Manager → "ComfyUI Neural Network Toolkit NNT"). Needs matplotlib, seaborn, and sklearn (the confusion matrix is built by sklearn in the eval node; all in requirements). Heavy install as always - torch, numpy, pandas, transformers, statsmodels, onnx, shap 0.41.0 - but this is one of the highest-value learning nodes in the pack: it turns "88% accuracy" into a picture of exactly which digits you're failing at.