Spaces:
Running
Running
| from collections import deque | |
| from dataclasses import dataclass | |
| import cv2 | |
| import numpy as np | |
| WIDTH = 320 # frames are downscaled to this width before processing | |
| # defaults; these can be changed live from the settings page | |
| VAR_THRESHOLD = 16 # MOG2 sensitivity; higher ignores subtler changes | |
| MIN_BLOB_FRAC = 0.0002 # smallest blob, as a fraction of the image, worth reacting to | |
| WARMUP = 15 # frames MOG2 needs to learn the background before its mask is useful | |
| # moving-camera detection | |
| FRAME_GAP = 6 # compare frames this far apart so slow motion still shows up | |
| DIFF_THRESHOLD = 20 # grayscale change, after alignment, that counts as motion | |
| MIN_FEATURES = 30 # tracked points needed to trust the alignment | |
| MIN_INLIER_FRAC = 0.5 # fraction of points that must agree on the camera motion | |
| MAX_MOTION_FRAC = 0.25 # more of the image than this "moving" means alignment failed | |
| KERNEL = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5)) | |
| BLOB_COLOR = (0, 255, 0) | |
| SMALL_BLOB_COLOR = (0, 0, 255) # largest blob, but below MIN_BLOB_FRAC | |
| LABEL_COLOR = (0, 200, 255) | |
| DEBUG_WIDTH = 1280 # panels are upscaled so labels and boxes stay crisp | |
| class Snapshot: | |
| """One frame's detection state, for the debug view.""" | |
| frame: np.ndarray # downscaled BGR frame | |
| raw: np.ndarray # detector output before cleaning | |
| raw_label: str | |
| mask: np.ndarray | None # cleaned mask; None until the detector is ready | |
| box: tuple[int, ...] | None # largest blob | |
| ok: bool # largest blob is big enough to track | |
| class MotionDetector: | |
| """Find the largest moving blob in a stream of frames from a still camera.""" | |
| def __init__(self) -> None: | |
| self.var_threshold = VAR_THRESHOLD | |
| self.min_blob_frac = MIN_BLOB_FRAC | |
| # replaced as a whole so readers on other threads see a consistent set | |
| self.snapshot: Snapshot | None = None | |
| self.reset() | |
| def set_var_threshold(self, value: float) -> None: | |
| """Change MOG2 sensitivity without forgetting the learned background.""" | |
| self.var_threshold = value | |
| self.sub.setVarThreshold(value) | |
| def reset(self) -> None: | |
| """Forget the learned background, e.g. after the camera moved.""" | |
| # shadow detection also flags anything moderately darker than the background | |
| # (e.g. gray clothes against a white wall), so treat every change as motion | |
| self.sub = cv2.createBackgroundSubtractorMOG2( | |
| varThreshold=self.var_threshold, detectShadows=False | |
| ) | |
| self.n = 0 | |
| def detect(self, frame: np.ndarray) -> tuple[float, float] | None: | |
| """Return the (x, y) full-res pixel center of the largest moving blob, or None.""" | |
| small, scale = shrink(frame) | |
| raw = self.sub.apply(small) | |
| self.n += 1 | |
| if self.n < WARMUP: | |
| self.snapshot = Snapshot(small, raw, "raw mask", None, None, False) | |
| return None | |
| self.snapshot, target = pick( | |
| small, raw, "raw mask", clean(raw), scale, self.min_blob_frac | |
| ) | |
| return target | |
| class ScanDetector: | |
| """Find the largest moving blob while the camera turns, by cancelling its motion. | |
| Consecutive frames are aligned with a homography fitted to tracked background | |
| features (a turning head is close to a pure rotation, which a homography models | |
| exactly), so after warping only independently moving things differ. | |
| """ | |
| def __init__(self) -> None: | |
| self.diff_threshold = DIFF_THRESHOLD | |
| self.min_blob_frac = MIN_BLOB_FRAC | |
| self.snapshot: Snapshot | None = None | |
| self.reset() | |
| def reset(self) -> None: | |
| """Drop frame history, e.g. when a new sweep starts.""" | |
| self.history: deque[np.ndarray] = deque(maxlen=FRAME_GAP + 1) | |
| def detect(self, frame: np.ndarray) -> tuple[float, float] | None: | |
| """Return the (x, y) full-res pixel center of the largest moving blob, or None.""" | |
| small, scale = shrink(frame) | |
| gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY) | |
| if self.history and np.array_equal(gray, self.history[-1]): | |
| return None # repeated frame; comparing it would hide motion | |
| self.history.append(gray) | |
| if len(self.history) <= FRAME_GAP: | |
| return None | |
| old = self.history[0] | |
| H = align(old, gray) | |
| if H is None: | |
| blank = np.zeros_like(gray) | |
| self.snapshot = Snapshot(small, blank, "no alignment", blank, None, False) | |
| return None | |
| h, w = gray.shape | |
| warped = cv2.warpPerspective(old, H, (w, h)) | |
| # pixels warped in from outside the old frame have no data; trim their edges | |
| valid = cv2.warpPerspective(np.full_like(old, 255), H, (w, h)) | |
| valid = cv2.erode(valid, KERNEL, iterations=2) > 0 | |
| diff = cv2.absdiff(gray, warped) | |
| # auto-exposure shifts the whole image while panning; remove that offset | |
| diff = cv2.subtract(diff, int(np.median(diff[valid]))) | |
| diff[~valid] = 0 | |
| _, mask = cv2.threshold(diff, self.diff_threshold, 255, cv2.THRESH_BINARY) | |
| mask = clean(mask) | |
| if mask.mean() > MAX_MOTION_FRAC * 255: | |
| mask[:] = 0 # everything "moving" means the alignment was off | |
| self.snapshot, target = pick( | |
| small, diff, "aligned difference", mask, scale, self.min_blob_frac | |
| ) | |
| return target | |
| def debug_image(snap: Snapshot | None) -> np.ndarray | None: | |
| """Render a snapshot's frame above its raw and cleaned masks.""" | |
| if snap is None: | |
| return None | |
| frame = upscale(snap.frame, DEBUG_WIDTH, cv2.INTER_LINEAR) | |
| if snap.box is not None: | |
| f = DEBUG_WIDTH / snap.frame.shape[1] | |
| x, y, w, h = (round(v * f) for v in snap.box) | |
| color = BLOB_COLOR if snap.ok else SMALL_BLOB_COLOR | |
| cv2.rectangle(frame, (x, y), (x + w, y + h), color, 3) | |
| # during warmup there is no cleaned mask yet | |
| mask = np.zeros_like(snap.raw) if snap.mask is None else snap.mask | |
| masks = [ | |
| label(upscale(cv2.cvtColor(m, cv2.COLOR_GRAY2BGR), DEBUG_WIDTH // 2), text) | |
| for text, m in ((snap.raw_label, snap.raw), ("cleaned mask", mask)) | |
| ] | |
| return np.vstack([label(frame, "camera"), np.hstack(masks)]) | |
| def shrink(frame: np.ndarray) -> tuple[np.ndarray, float]: | |
| """Downscale and blur a frame; return it with the factor back to full res.""" | |
| scale = frame.shape[1] / WIDTH | |
| small = cv2.resize( | |
| frame, None, fx=1 / scale, fy=1 / scale, interpolation=cv2.INTER_AREA | |
| ) | |
| return cv2.GaussianBlur(small, (5, 5), 0), scale | |
| def clean(mask: np.ndarray) -> np.ndarray: | |
| """Drop speckle, then join fragments of the same object.""" | |
| mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, KERNEL) | |
| return cv2.morphologyEx(mask, cv2.MORPH_CLOSE, KERNEL, iterations=2) | |
| def pick( | |
| small: np.ndarray, | |
| raw: np.ndarray, | |
| raw_label: str, | |
| mask: np.ndarray, | |
| scale: float, | |
| min_blob_frac: float, | |
| ) -> tuple[Snapshot, tuple[float, float] | None]: | |
| """Choose the largest blob in a cleaned mask; return a snapshot and its target.""" | |
| blob = largest_blob(mask) | |
| ok = blob is not None and blob[2] >= min_blob_frac * mask.size | |
| snap = Snapshot(small, raw, raw_label, mask, blob and blob[0], ok) | |
| if not ok: | |
| return snap, None | |
| x, y = blob[1] | |
| return snap, (float(x * scale), float(y * scale)) | |
| def align(old: np.ndarray, new: np.ndarray) -> np.ndarray | None: | |
| """Estimate the homography mapping old onto new from tracked features.""" | |
| pts = cv2.goodFeaturesToTrack(old, maxCorners=300, qualityLevel=0.01, minDistance=8) | |
| if pts is None or len(pts) < MIN_FEATURES: | |
| return None | |
| nxt, status, _ = cv2.calcOpticalFlowPyrLK(old, new, pts, None) | |
| found = status.ravel() == 1 | |
| if found.sum() < MIN_FEATURES: | |
| return None | |
| # RANSAC ignores points on moving objects, but a bad fit can still slip through | |
| H, inliers = cv2.findHomography(pts[found], nxt[found], cv2.RANSAC, 2.0) | |
| if H is None or inliers.mean() < MIN_INLIER_FRAC: | |
| return None | |
| return H | |
| def upscale(img: np.ndarray, width: int, interp: int = cv2.INTER_NEAREST) -> np.ndarray: | |
| """Enlarge an image to the given width; nearest-neighbor keeps mask pixels sharp.""" | |
| f = width / img.shape[1] | |
| return cv2.resize(img, None, fx=f, fy=f, interpolation=interp) | |
| def label(img: np.ndarray, text: str) -> np.ndarray: | |
| """Draw an outlined caption in the top-left corner.""" | |
| for color, thickness in (((0, 0, 0), 4), (LABEL_COLOR, 1)): | |
| cv2.putText( | |
| img, | |
| text, | |
| (8, 24), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| 0.7, | |
| color, | |
| thickness, | |
| cv2.LINE_AA, | |
| ) | |
| return img | |
| def largest_blob( | |
| mask: np.ndarray, | |
| ) -> tuple[tuple[int, ...], tuple[float, float], int] | None: | |
| """Return the (x, y, w, h) box, centroid and area of the largest blob.""" | |
| n, _, stats, centroids = cv2.connectedComponentsWithStats(mask) | |
| if n < 2: | |
| return None | |
| # label 0 is the background | |
| i = 1 + int(np.argmax(stats[1:, cv2.CC_STAT_AREA])) | |
| box = tuple(int(v) for v in stats[i, :4]) | |
| return box, tuple(centroids[i]), int(stats[i, cv2.CC_STAT_AREA]) | |