72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
"""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()
|