Calculate and plot confusion matrix to a SummaryWriter. Args: writer (SummaryWriter): the SummaryWriter to write the matrix to. cmtx (ndarray): confusion matrix. num_classes (int): total number of classes. global_step (Optional[int]): current step. su
(
writer,
cmtx,
num_classes,
global_step=None,
subset_ids=None,
class_names=None,
tag="Confusion Matrix",
figsize=None,
)
| 236 | |
| 237 | |
| 238 | def add_confusion_matrix( |
| 239 | writer, |
| 240 | cmtx, |
| 241 | num_classes, |
| 242 | global_step=None, |
| 243 | subset_ids=None, |
| 244 | class_names=None, |
| 245 | tag="Confusion Matrix", |
| 246 | figsize=None, |
| 247 | ): |
| 248 | """ |
| 249 | Calculate and plot confusion matrix to a SummaryWriter. |
| 250 | Args: |
| 251 | writer (SummaryWriter): the SummaryWriter to write the matrix to. |
| 252 | cmtx (ndarray): confusion matrix. |
| 253 | num_classes (int): total number of classes. |
| 254 | global_step (Optional[int]): current step. |
| 255 | subset_ids (list of ints): a list of label indices to keep. |
| 256 | class_names (list of strs, optional): a list of all class names. |
| 257 | tag (str or list of strs): name(s) of the confusion matrix image. |
| 258 | figsize (Optional[float, float]): the figure size of the confusion matrix. |
| 259 | If None, default to [6.4, 4.8]. |
| 260 | |
| 261 | """ |
| 262 | if subset_ids is None or len(subset_ids) != 0: |
| 263 | # If class names are not provided, use class indices as class names. |
| 264 | if class_names is None: |
| 265 | class_names = [str(i) for i in range(num_classes)] |
| 266 | # If subset is not provided, take every classes. |
| 267 | if subset_ids is None: |
| 268 | subset_ids = list(range(num_classes)) |
| 269 | |
| 270 | sub_cmtx = cmtx[subset_ids, :][:, subset_ids] |
| 271 | sub_names = [class_names[j] for j in subset_ids] |
| 272 | |
| 273 | sub_cmtx = vis_utils.plot_confusion_matrix( |
| 274 | sub_cmtx, |
| 275 | num_classes=len(subset_ids), |
| 276 | class_names=sub_names, |
| 277 | figsize=figsize, |
| 278 | ) |
| 279 | # Add the confusion matrix image to writer. |
| 280 | writer.add_figure(tag=tag, figure=sub_cmtx, global_step=global_step) |
| 281 | |
| 282 | |
| 283 | def plot_hist( |