| """Visualization and evaluation helpers used by the inference entrypoints.""" |
|
|
| import matplotlib |
| matplotlib.use("Agg") |
|
|
| import cv2 |
| import matplotlib.pyplot as plt |
| import numpy as np |
| from sklearn.metrics import confusion_matrix |
|
|
|
|
| def get_cmap(name, n): |
| """Discrete colormap, compatible with matplotlib >= 3.9. |
| |
| ``plt.cm.get_cmap`` was removed in 3.9; ``plt.get_cmap(...).resampled(n)`` |
| is the supported spelling, with a fallback for older releases. |
| """ |
| cmap = plt.get_cmap(name) |
| return cmap.resampled(n) if hasattr(cmap, "resampled") else plt.cm.get_cmap(name, n) |
|
|
|
|
| def quick_resize(img, max_dim=1024, is_mask=False): |
| """Shrink so the longest edge is <= max_dim, preserving aspect ratio. |
| |
| Masks use nearest-neighbour so class indices survive; images/probability |
| maps use area interpolation. |
| """ |
| h, w = img.shape[:2] |
| scale = max_dim / float(max(h, w)) |
| if scale < 1.0: |
| new_w, new_h = int(w * scale), int(h * scale) |
| interp = cv2.INTER_NEAREST if is_mask else cv2.INTER_AREA |
| img = cv2.resize(img, (new_w, new_h), interpolation=interp) |
| return img |
|
|
|
|
| def vis_img(image, pred_mask, foreground_probs_all, template_all, save_path, max_dim=1024): |
| """One row per class (image | probability map) plus a final merged-mask row.""" |
| img_small = quick_resize(image, max_dim=max_dim) |
| pred_mask_small = quick_resize(pred_mask, max_dim=max_dim, is_mask=True) |
| probs_small = [quick_resize(p, max_dim=max_dim) for p in foreground_probs_all] |
|
|
| num_rows = len(probs_small) + 1 |
| fig, axs = plt.subplots(num_rows, 2, figsize=(12, 4 * num_rows), squeeze=False) |
|
|
| for i in range(num_rows - 1): |
| axs[i, 0].imshow(img_small) |
| axs[i, 0].set_title(template_all[i]) |
| axs[i, 0].axis("off") |
|
|
| axs[i, 1].imshow(probs_small[i], cmap="gray", vmin=0, vmax=1) |
| axs[i, 1].set_title("Predicted Probability") |
| axs[i, 1].axis("off") |
|
|
| n_classes = len(template_all) + 1 |
| cmap = get_cmap("tab20", n_classes) |
|
|
| axs[num_rows - 1, 0].imshow(img_small) |
| axs[num_rows - 1, 0].set_title("Raw Image") |
| axs[num_rows - 1, 0].axis("off") |
|
|
| im1 = axs[num_rows - 1, 1].imshow(pred_mask_small, cmap=cmap, vmin=0, |
| vmax=n_classes - 1, interpolation="nearest") |
| axs[num_rows - 1, 1].set_title("Predicted Mask") |
| axs[num_rows - 1, 1].axis("off") |
|
|
| cbar = fig.colorbar(im1, ax=axs[num_rows - 1, 1], orientation="vertical", |
| fraction=0.02, pad=0.04) |
| cbar.set_ticks(range(n_classes)) |
| cbar.set_ticklabels(["background"] + list(template_all)) |
|
|
| fig.savefig(save_path, bbox_inches="tight", dpi=200) |
| plt.close(fig) |
|
|
|
|
| def vis_img_bbx(image, pred_mask, true_mask, foreground_probs_all, gt_masks_all, |
| template_all, save_path, max_dim=1024): |
| """One row per class (image | prediction | ground truth) plus a merged row.""" |
| img_small = quick_resize(image, max_dim=max_dim) |
| pred_mask_small = quick_resize(pred_mask, max_dim=max_dim, is_mask=True) |
| true_mask_small = quick_resize(true_mask, max_dim=max_dim, is_mask=True) |
| probs_small = [quick_resize(p, max_dim=max_dim) for p in foreground_probs_all] |
| gts_small = [quick_resize(g, max_dim=max_dim, is_mask=True) for g in gt_masks_all] |
|
|
| num_rows = len(probs_small) + 1 |
| fig, axs = plt.subplots(num_rows, 3, figsize=(12, 4 * num_rows), squeeze=False) |
|
|
| for i in range(num_rows - 1): |
| axs[i, 0].imshow(img_small) |
| axs[i, 0].set_title(template_all[i]) |
| axs[i, 0].axis("off") |
|
|
| axs[i, 1].imshow(probs_small[i], cmap="gray", vmin=0, vmax=1) |
| axs[i, 1].set_title("Predicted Probability") |
| axs[i, 1].axis("off") |
|
|
| axs[i, 2].imshow(gts_small[i], cmap="gray", vmin=0, vmax=1) |
| axs[i, 2].set_title("Ground Truth Mask") |
| axs[i, 2].axis("off") |
|
|
| n_classes = len(template_all) + 1 |
| cmap = get_cmap("tab20", n_classes) |
|
|
| axs[num_rows - 1, 0].imshow(img_small) |
| axs[num_rows - 1, 0].set_title("Raw Image") |
| axs[num_rows - 1, 0].axis("off") |
|
|
| axs[num_rows - 1, 1].imshow(pred_mask_small, cmap=cmap, vmin=0, |
| vmax=n_classes - 1, interpolation="nearest") |
| axs[num_rows - 1, 1].set_title("Predicted Mask") |
| axs[num_rows - 1, 1].axis("off") |
|
|
| im1 = axs[num_rows - 1, 2].imshow(true_mask_small, cmap=cmap, vmin=0, |
| vmax=n_classes - 1, interpolation="nearest") |
| axs[num_rows - 1, 2].set_title("Ground Truth Mask") |
| axs[num_rows - 1, 2].axis("off") |
|
|
| cbar = fig.colorbar(im1, ax=axs[num_rows - 1, 2], orientation="vertical", |
| fraction=0.02, pad=0.04) |
| cbar.set_ticks(range(n_classes)) |
| cbar.set_ticklabels(["background"] + list(template_all)) |
|
|
| fig.savefig(save_path, bbox_inches="tight", dpi=200) |
| plt.close(fig) |
|
|
|
|
| def save_prob_maps(foreground_probs_all, template_all, save_path="prob_maps.npz"): |
| """Save one probability map per class into a single .npz keyed by class name.""" |
| assert len(foreground_probs_all) == len(template_all), \ |
| "probability maps and class names must be the same length" |
|
|
| prob_dict = {name: prob.astype(np.float16) |
| for name, prob in zip(template_all, foreground_probs_all)} |
| np.savez_compressed(save_path, **prob_dict) |
| print(f"Saved probability maps to {save_path}") |
|
|
|
|
| def evaluate_segmentation(pred_mask, true_mask, eps=1e-7): |
| """Binary segmentation metrics: (dice, accuracy, precision, recall).""" |
| pred_bin = (pred_mask > 0).astype(np.float32).reshape(-1) |
| true_bin = (true_mask > 0).astype(np.float32).reshape(-1) |
|
|
| intersection = np.sum(pred_bin * true_bin) |
| union = np.sum(pred_bin) + np.sum(true_bin) |
|
|
| dice = (2.0 * intersection + eps) / (union + eps) |
| acc = np.mean(pred_bin == true_bin) |
|
|
| tp = np.sum(pred_bin * true_bin) |
| fp = np.sum(pred_bin * (1.0 - true_bin)) |
| fn = np.sum((1.0 - pred_bin) * true_bin) |
|
|
| precision = tp / (tp + fp + eps) |
| recall = tp / (tp + fn + eps) |
|
|
| return dice, acc, precision, recall |
|
|
|
|
| def compute_multi_class_metrics(gt, pred): |
| """Macro-averaged IoU / Dice / precision / recall over the classes present in gt.""" |
| metrics = {} |
| epsilon = 1e-7 |
|
|
| num_classes = max(int(np.amax(gt)), int(np.amax(pred))) + 1 |
| cm = confusion_matrix(gt.flatten(), pred.flatten(), labels=np.arange(num_classes)) |
|
|
| IoU, Dice, Precision, Recall = [], [], [], [] |
|
|
| for i in range(num_classes): |
| TP = cm[i, i] |
| FP = cm[:, i].sum() - TP |
| FN = cm[i, :].sum() - TP |
|
|
| |
| if (TP + FN) == 0: |
| continue |
|
|
| IoU.append(TP / (TP + FP + FN + epsilon)) |
| Dice.append(2 * TP / (2 * TP + FP + FN + epsilon)) |
| Precision.append(TP / (TP + FP + epsilon)) |
| Recall.append(TP / (TP + FN + epsilon)) |
|
|
| metrics["IoU"] = np.mean(IoU) |
| metrics["Dice"] = np.mean(Dice) |
| metrics["Precision"] = np.mean(Precision) |
| metrics["Recall"] = np.mean(Recall) |
|
|
| return metrics |
|
|
