move_detect/edge/snapshot.py

45 lines
1.5 KiB
Python

from collections import deque
import cv2
from edge.models import BBox
class RingBuffer:
def __init__(self, maxlen: int = 90):
self.frames: deque = deque(maxlen=maxlen)
def push(self, t: float, frame) -> None:
self.frames.append((t, frame))
def window(self, pre: int, post: int, index: int) -> list:
lo = max(0, index - pre)
hi = min(len(self.frames), index + post + 1)
return [f for (_t, f) in list(self.frames)[lo:hi]]
class FaceBlurrer:
def blur(self, frame, head_boxes: list[BBox]):
for b in head_boxes:
x1, y1 = int(max(0, b.x1)), int(max(0, b.y1))
x2, y2 = int(min(frame.shape[1], b.x2)), int(min(frame.shape[0], b.y2))
if x2 > x1 and y2 > y1:
roi = frame[y1:y2, x1:x2]
frame[y1:y2, x1:x2] = cv2.GaussianBlur(roi, (0, 0), sigmaX=12)
return frame
class Snapshotter:
def __init__(self, blurrer: FaceBlurrer, quality: int = 80):
self.blurrer = blurrer
self.quality = quality
def capture(self, frames: list, head_boxes_per_frame: list) -> list[bytes]:
out: list[bytes] = []
for frame, heads in list(zip(frames, head_boxes_per_frame))[:3]:
blurred = self.blurrer.blur(frame.copy(), heads)
ok, buf = cv2.imencode(".jpg", blurred,
[cv2.IMWRITE_JPEG_QUALITY, self.quality])
if ok:
out.append(buf.tobytes())
return out