Spaces:
Running on Zero
Running on Zero
File size: 5,595 Bytes
f234112 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """Visualization helpers for accident detection."""
from __future__ import annotations
import io
from pathlib import Path
from typing import Any
import matplotlib.patches as patches
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
from . import CLASS_ID_ACCIDENT, CLASS_ID_NON_ACCIDENT, CLASS_NAMES
PALETTE = {
CLASS_ID_ACCIDENT: (1.0, 0.0, 0.0), # red
CLASS_ID_NON_ACCIDENT: (0.0, 0.5, 1.0), # blue
}
def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
"""Convert a (3, H, W) float tensor in [0, 1] to a uint8 RGB array."""
arr = tensor.detach().cpu().numpy()
arr = np.transpose(arr, (1, 2, 0))
arr = np.clip(arr * 255.0, 0, 255).astype(np.uint8)
return arr
def draw_boxes(
image: np.ndarray,
boxes: np.ndarray,
labels: np.ndarray,
confidences: np.ndarray | None = None,
title: str = "",
line_width: int = 2,
font_size: int = 10,
) -> plt.Figure:
"""Draw bounding boxes on an RGB image array."""
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
ax.imshow(image)
ax.axis("off")
if title:
ax.set_title(title)
boxes = np.asarray(boxes)
labels = np.asarray(labels)
if boxes.size == 0:
return fig
for i in range(len(boxes)):
x1, y1, x2, y2 = boxes[i]
w = x2 - x1
h = y2 - y1
cat_id = int(labels[i])
color = PALETTE.get(cat_id, (0.5, 0.5, 0.5))
rect = patches.Rectangle(
(x1, y1), w, h, linewidth=line_width, edgecolor=color, facecolor="none"
)
ax.add_patch(rect)
label_text = CLASS_NAMES.get(cat_id, str(cat_id))
if confidences is not None:
label_text += f" {confidences[i]:.2f}"
ax.text(
x1, y1 - 4, label_text, color="white", fontsize=font_size,
bbox=dict(facecolor=color, alpha=0.7, edgecolor="none", pad=1),
)
plt.tight_layout()
return fig
def save_ground_truth_example(
item: dict[str, Any],
save_path: Path,
title: str = "Ground Truth",
) -> None:
"""Visualize one dataset item and save to disk."""
image = denormalize_image(item["image"])
boxes = item["boxes"].numpy()
labels = item["labels"].numpy()
fig = draw_boxes(image, boxes, labels, title=title)
save_path.parent.mkdir(parents=True, exist_ok=True)
fig.savefig(save_path, bbox_inches="tight", pad_inches=0.05)
plt.close(fig)
def save_prediction_example(
image_tensor: torch.Tensor,
gt_boxes: np.ndarray,
gt_labels: np.ndarray,
pred_boxes: np.ndarray,
pred_labels: np.ndarray,
pred_scores: np.ndarray,
save_path: Path,
title: str = "Predictions",
) -> None:
"""Visualize ground truth (dashed) and predictions (solid) on one image."""
fig, ax = plt.subplots(1, 1, figsize=(6, 6))
image = denormalize_image(image_tensor)
ax.imshow(image)
ax.axis("off")
if title:
ax.set_title(title)
# Ground truth boxes (dashed)
for box, label in zip(gt_boxes, gt_labels):
x1, y1, x2, y2 = box
color = PALETTE.get(int(label), (0.5, 0.5, 0.5))
rect = patches.Rectangle(
(x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor=color,
facecolor="none", linestyle="--", alpha=0.7,
)
ax.add_patch(rect)
ax.text(
x1, y2 + 12, f"GT {CLASS_NAMES.get(int(label), label)}",
color="white", fontsize=9,
bbox=dict(facecolor=color, alpha=0.5, edgecolor="none", pad=1),
)
# Predicted boxes (solid)
for box, label, score in zip(pred_boxes, pred_labels, pred_scores):
x1, y1, x2, y2 = box
color = PALETTE.get(int(label), (0.5, 0.5, 0.5))
rect = patches.Rectangle(
(x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor=color,
facecolor="none",
)
ax.add_patch(rect)
ax.text(
x1, y1 - 4, f"Pred {CLASS_NAMES.get(int(label), label)} {score:.2f}",
color="white", fontsize=9,
bbox=dict(facecolor=color, alpha=0.7, edgecolor="none", pad=1),
)
save_path.parent.mkdir(parents=True, exist_ok=True)
plt.tight_layout()
fig.savefig(save_path, bbox_inches="tight", pad_inches=0.05)
plt.close(fig)
def save_grid(
items: list[dict[str, Any]],
save_path: Path,
title: str = "Ground Truth Examples",
max_cols: int = 5,
) -> None:
"""Save a grid of ground-truth examples."""
n = len(items)
cols = min(n, max_cols)
rows = (n + cols - 1) // cols
fig, axes = plt.subplots(rows, cols, figsize=(3 * cols, 3 * rows))
if n == 1:
axes = np.array([axes])
axes = axes.flatten()
for ax, item in zip(axes, items):
image = denormalize_image(item["image"])
ax.imshow(image)
boxes = item["boxes"].numpy()
labels = item["labels"].numpy()
for box, label in zip(boxes, labels):
x1, y1, x2, y2 = box
color = PALETTE.get(int(label), (0.5, 0.5, 0.5))
rect = patches.Rectangle(
(x1, y1), x2 - x1, y2 - y1, linewidth=2, edgecolor=color,
facecolor="none",
)
ax.add_patch(rect)
ax.axis("off")
ax.set_title(f"id={item['image_id']} n={item['num_objects']}", fontsize=8)
for ax in axes[n:]:
ax.axis("off")
fig.suptitle(title, fontsize=12)
save_path.parent.mkdir(parents=True, exist_ok=True)
plt.tight_layout()
fig.savefig(save_path, bbox_inches="tight", pad_inches=0.1)
plt.close(fig)
|