27 lines
807 B
Python
27 lines
807 B
Python
"""Build a binary mask of the censored regions from detections."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from ..detection.types import Detection
|
|
|
|
|
|
def detections_to_mask(
|
|
shape: tuple[int, int], detections: list[Detection], dilate: int = 0
|
|
) -> np.ndarray:
|
|
"""White (255) over every detected region (polygon if present, else bbox)."""
|
|
h, w = shape[:2]
|
|
mask = np.zeros((h, w), np.uint8)
|
|
for d in detections:
|
|
if len(d.polygon) >= 3:
|
|
cv2.fillPoly(mask, [np.array(d.polygon, np.int32)], 255)
|
|
else:
|
|
x, y, bw, bh = d.bbox
|
|
cv2.rectangle(mask, (x, y), (x + bw, y + bh), 255, -1)
|
|
if dilate > 0:
|
|
k = np.ones((dilate * 2 + 1, dilate * 2 + 1), np.uint8)
|
|
mask = cv2.dilate(mask, k)
|
|
return mask
|