feat(Q2): heatmap generation with bbox accumulation and colormap
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
parent
7ccd818a93
commit
66a53f0cd8
17
tests/unit/test_heatmap.py
Normal file
17
tests/unit/test_heatmap.py
Normal file
@ -0,0 +1,17 @@
|
||||
import numpy as np
|
||||
from vigilar.detection.heatmap import accumulate_detections, render_heatmap_png
|
||||
|
||||
def test_accumulate_detections_returns_grid():
|
||||
bboxes = [[0.5, 0.5, 0.1, 0.1], [0.5, 0.5, 0.1, 0.1], [0.2, 0.3, 0.1, 0.1]]
|
||||
grid = accumulate_detections(bboxes, grid_w=64, grid_h=36)
|
||||
assert grid.shape == (36, 64)
|
||||
assert grid.sum() > 0
|
||||
assert grid[18, 32] > grid[0, 0] # center of [0.5, 0.5]
|
||||
|
||||
def test_generate_heatmap_returns_png_bytes():
|
||||
bboxes = [[0.3, 0.4, 0.1, 0.1]] * 20
|
||||
grid = accumulate_detections(bboxes)
|
||||
frame = np.zeros((360, 640, 3), dtype=np.uint8)
|
||||
png_bytes = render_heatmap_png(grid, frame)
|
||||
assert isinstance(png_bytes, bytes)
|
||||
assert png_bytes[:4] == b"\x89PNG"
|
||||
71
vigilar/detection/heatmap.py
Normal file
71
vigilar/detection/heatmap.py
Normal file
@ -0,0 +1,71 @@
|
||||
"""Activity heatmap generation from detection bounding boxes."""
|
||||
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
GRID_W = 64
|
||||
GRID_H = 36
|
||||
|
||||
|
||||
def accumulate_detections(
|
||||
bboxes: list[list[float]], grid_w: int = GRID_W, grid_h: int = GRID_H,
|
||||
) -> np.ndarray:
|
||||
grid = np.zeros((grid_h, grid_w), dtype=np.float32)
|
||||
for bbox in bboxes:
|
||||
if len(bbox) < 4:
|
||||
continue
|
||||
x, y, w, h = bbox
|
||||
cx = x
|
||||
cy = y
|
||||
gx = max(0, min(grid_w - 1, int(cx * grid_w)))
|
||||
gy = max(0, min(grid_h - 1, int(cy * grid_h)))
|
||||
grid[gy, gx] += 1
|
||||
return grid
|
||||
|
||||
|
||||
def _gaussian_blur(grid: np.ndarray, sigma: float = 2.0) -> np.ndarray:
|
||||
size = int(sigma * 3) * 2 + 1
|
||||
x = np.arange(size) - size // 2
|
||||
kernel_1d = np.exp(-x ** 2 / (2 * sigma ** 2))
|
||||
kernel_1d = kernel_1d / kernel_1d.sum()
|
||||
blurred = np.apply_along_axis(lambda row: np.convolve(row, kernel_1d, mode="same"), 1, grid)
|
||||
blurred = np.apply_along_axis(lambda col: np.convolve(col, kernel_1d, mode="same"), 0, blurred)
|
||||
return blurred
|
||||
|
||||
|
||||
def _apply_colormap(normalized: np.ndarray) -> np.ndarray:
|
||||
h, w = normalized.shape
|
||||
rgb = np.zeros((h, w, 3), dtype=np.uint8)
|
||||
low = normalized <= 0.5
|
||||
high = ~low
|
||||
t_low = normalized[low] * 2
|
||||
rgb[low, 0] = (t_low * 255).astype(np.uint8)
|
||||
rgb[low, 1] = (t_low * 255).astype(np.uint8)
|
||||
rgb[low, 2] = ((1 - t_low) * 255).astype(np.uint8)
|
||||
t_high = (normalized[high] - 0.5) * 2
|
||||
rgb[high, 0] = np.full(t_high.shape, 255, dtype=np.uint8)
|
||||
rgb[high, 1] = ((1 - t_high) * 255).astype(np.uint8)
|
||||
rgb[high, 2] = np.zeros(t_high.shape, dtype=np.uint8)
|
||||
return rgb
|
||||
|
||||
|
||||
def render_heatmap_png(grid: np.ndarray, frame: np.ndarray, alpha: float = 0.5) -> bytes:
|
||||
from PIL import Image
|
||||
blurred = _gaussian_blur(grid)
|
||||
max_val = blurred.max()
|
||||
normalized = blurred / max_val if max_val > 0 else blurred
|
||||
heatmap_rgb = _apply_colormap(normalized)
|
||||
frame_h, frame_w = frame.shape[:2]
|
||||
heatmap_img = Image.fromarray(heatmap_rgb).resize((frame_w, frame_h), Image.BILINEAR)
|
||||
frame_img = Image.fromarray(frame[:, :, ::-1]) # BGR to RGB
|
||||
blended = Image.blend(frame_img, heatmap_img, alpha=alpha)
|
||||
buf = io.BytesIO()
|
||||
blended.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
Loading…
Reference in New Issue
Block a user