18 lines
720 B
Python
18 lines
720 B
Python
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"
|