code/blur_video.py

#!/usr/bin/env python3
"""Detect, track and pixelate (or black out) faces in a video. Audio is preserved via ffmpeg.

Detection runs every N frames; between detections ByteTrack's Kalman filter carries the boxes forward.
Tuned for recall (privacy): low confidence threshold, box padding, raw detections are blurred immediately
(without waiting for track confirmation) and recently lost tracks are held for a few frames.

--keep: leave ONE person's face visible and blur everybody else. Give one or more reference photos of that
person. Fail-safe by design: a face stays blurred unless its track was confirmed as the target by several
consecutive face-recognition matches; a clear non-match, a track box jump, or going too long without a
successful re-check puts the blur back.

Examples:
  python blur_video.py in.mp4 out.mp4 --model export/face_yolo11n_fp16.onnx -n 1
  python blur_video.py in.mp4 out.mp4 --mode black
  python blur_video.py in.mp4 out.mp4 -n 1 --keep me1.jpg me2.jpg
"""
import argparse
import queue
import shutil
import subprocess
import sys
import threading
import time
from types import SimpleNamespace

import cv2
import numpy as np
from ultralytics import YOLO
from ultralytics.trackers.byte_tracker import BYTETracker


def pad_box(x1, y1, x2, y2, pad, w, h):
    bw, bh = x2 - x1, y2 - y1
    x1, x2 = x1 - bw * pad, x2 + bw * pad
    y1, y2 = y1 - bh * pad, y2 + bh * pad
    return max(0, int(np.floor(x1))), max(0, int(np.floor(y1))), min(w, int(np.ceil(x2))), min(h, int(np.ceil(y2)))


def iou(a, b):
    ix = max(0.0, min(a[2], b[2]) - max(a[0], b[0]))
    iy = max(0.0, min(a[3], b[3]) - max(a[1], b[1]))
    inter = ix * iy
    union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
    return inter / union if union > 0 else 0.0


def pixelate(frame, box, blocks):
    """Downscale the region so its shorter side has `blocks` cells, then nearest-neighbour upscale."""
    x1, y1, x2, y2 = box
    if x2 - x1 < 2 or y2 - y1 < 2:
        return
    roi = frame[y1:y2, x1:x2]
    h, w = roi.shape[:2]
    cell = max(1.0, min(w, h) / blocks)
    small = cv2.resize(roi, (max(1, round(w / cell)), max(1, round(h / cell))), interpolation=cv2.INTER_AREA)
    frame[y1:y2, x1:x2] = cv2.resize(small, (w, h), interpolation=cv2.INTER_NEAREST)


def blackout(frame, box):
    x1, y1, x2, y2 = box
    frame[y1:y2, x1:x2] = 0


def open_ffmpeg_writer(src, dst, w, h, fps, crf):
    cmd = [
        "ffmpeg", "-y", "-loglevel", "error",
        "-f", "rawvideo", "-pix_fmt", "bgr24", "-s", f"{w}x{h}", "-r", f"{fps}", "-i", "-",
        "-i", src,
        "-map", "0:v:0", "-map", "1:a?",
        "-c:v", "libx264", "-preset", "fast", "-crf", str(crf), "-pix_fmt", "yuv420p",
        "-c:a", "copy", "-shortest", "-movflags", "+faststart", dst,
    ]
    return subprocess.Popen(cmd, stdin=subprocess.PIPE)


class KeepTracker:
    """Decides which ByteTrack ids are the person to keep visible. Everything not confirmed stays blurred."""

    def __init__(self, fid, refs, thr, min_size, hits_needed, recheck, expire, budget):
        self.fid, self.refs, self.thr, self.min_size = fid, refs, thr, min_size
        self.hits_needed, self.recheck, self.expire, self.budget = hits_needed, recheck, expire, budget
        self.state = {}  # tid -> dict(hits, kept, last_ok, last_try, box)
        self.seen, self.ever_kept = set(), set()

    def is_kept(self, tid, idx):
        s = self.state.get(tid)
        if s is None or not s["kept"]:
            return False
        if idx - s["last_ok"] > self.expire:  # could not re-confirm for too long -> blur again
            s["kept"], s["hits"] = False, 0
            return False
        return True

    def update(self, frame, idx, tracked):
        """tracked: list of (xyxy, tid) from a detection frame (clean, un-blurred frame)."""
        todo = []
        for box, tid in tracked:
            self.seen.add(tid)
            s = self.state.setdefault(tid, dict(hits=0, kept=False, last_ok=-10**9, last_try=-10**9, box=box))
            if s["kept"] and iou(s["box"], box) < 0.3:  # box jumped: possible ID switch -> distrust
                s["kept"], s["hits"] = False, 0
            s["box"] = box
            if not self.is_kept(tid, idx):
                todo.append((0, box, tid))  # undecided first
            elif idx - s["last_try"] >= self.recheck:
                todo.append((1, box, tid))
        todo.sort(key=lambda t: (t[0], self.state[t[2]]["last_try"]))
        for _, box, tid in todo[:self.budget]:
            s = self.state[tid]
            s["last_try"] = idx
            e = self.fid.embed_box(frame, box, min_size=self.min_size)
            if e is None:
                continue  # inconclusive: keep old state (a kept track lapses via `expire`)
            if self.fid.similarity(e, self.refs) >= self.thr:
                s["hits"] += 1
                s["last_ok"] = idx
                if s["hits"] >= self.hits_needed:
                    s["kept"] = True
                    self.ever_kept.add(tid)
            else:  # clear non-match: blur immediately
                s["hits"], s["kept"] = 0, False


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("input")
    ap.add_argument("output")
    ap.add_argument("--model", default="runs/face_yolo11n/weights/best.pt", help=".pt or .onnx face detector")
    ap.add_argument("-n", "--every", type=int, default=3, help="run detector every N frames (default 3)")
    ap.add_argument("--conf", type=float, default=0.25, help="detection confidence threshold (default 0.25)")
    ap.add_argument("--pad", type=float, default=0.15, help="pad each box by this fraction per side (default 0.15)")
    ap.add_argument("--mode", choices=["pixelate", "black"], default="pixelate",
                    help="pixelate faces (default) or fill the padded box with solid black")
    ap.add_argument("--blocks", type=int, default=8,
                    help="pixel blocks across a face's shorter side; LOWER = stronger pixelation (default 8)")
    ap.add_argument("--hold", type=int, default=8, help="keep blurring a lost track's prediction this many frames")
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--device", default=None, help="e.g. 0 or cpu (default: auto)")
    ap.add_argument("--crf", type=int, default=18, help="x264 quality (lower = better)")
    k = ap.add_argument_group("keep one person visible (face recognition)")
    k.add_argument("--keep", nargs="+", metavar="REF.jpg",
                   help="reference photo(s) of the person to leave unblurred; everyone else is blurred")
    k.add_argument("--keep-threshold", type=float, default=0.45,
                   help="cosine similarity needed to count as the same person; HIGHER = stricter (default 0.45)")
    k.add_argument("--keep-hits", type=int, default=2, help="consecutive matches before a track is unblurred")
    k.add_argument("--keep-min-size", type=int, default=40,
                   help="faces smaller than this many px are never recognised, so they stay blurred (default 40)")
    k.add_argument("--keep-recheck", type=int, default=8, help="re-verify a kept track every this many frames")
    k.add_argument("--keep-expire", type=int, default=24,
                   help="blur a kept track again if it could not be re-confirmed for this many frames")
    k.add_argument("--keep-budget", type=int, default=4, help="max recognitions per frame (speed cap)")
    a = ap.parse_args()

    if shutil.which("ffmpeg") is None:
        sys.exit("ffmpeg not found on PATH")
    if a.every < 1 or a.blocks < 1:
        sys.exit("--every and --blocks must be >= 1")

    keeper = None
    if a.keep:
        from face_id import FaceID
        fid = FaceID("cpu" if a.device == "cpu" else "0")
        refs = fid.reference_embeddings(a.keep)
        print(f"keeping the person in {len(refs)} reference photo(s); threshold {a.keep_threshold}")
        keeper = KeepTracker(fid, refs, a.keep_threshold, a.keep_min_size, a.keep_hits,
                             a.keep_recheck, a.keep_expire, a.keep_budget)

    cap = cv2.VideoCapture(a.input)
    if not cap.isOpened():
        sys.exit(f"cannot open {a.input}")
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
    ok, frame = cap.read()
    if not ok:
        sys.exit("no frames in input")
    H, W = frame.shape[:2]

    model = YOLO(a.model, task="detect")
    tracker = BYTETracker(SimpleNamespace(
        track_high_thresh=a.conf, track_low_thresh=0.1, new_track_thresh=a.conf,
        track_buffer=max(30, 4 * a.every), match_thresh=0.8, fuse_score=True))
    writer = open_ffmpeg_writer(a.input, a.output, W, H, fps, a.crf)

    # Decode and encode run in their own threads so they overlap with detection (frame order is preserved:
    # one reader, one writer, FIFO queues). cv2/ffmpeg pipe I/O release the GIL.
    rq, wq, write_failed = queue.Queue(maxsize=48), queue.Queue(maxsize=48), []
    first_frame = frame

    def reader():
        rq.put(first_frame)
        while True:
            okr, f = cap.read()
            if not okr:
                rq.put(None)
                return
            rq.put(f)

    def writer_loop():
        try:
            while True:
                f = wq.get()
                if f is None:
                    return
                writer.stdin.write(f.tobytes())
        except (BrokenPipeError, OSError):
            write_failed.append(True)
            while wq.get() is not None:  # drain so the main thread never blocks on a full queue
                pass

    threading.Thread(target=reader, daemon=True).start()
    wthread = threading.Thread(target=writer_loop)
    wthread.start()

    t0 = time.perf_counter()
    idx = n_det = 0
    while True:
        frame = rq.get()
        if frame is None:
            break
        items = []  # (xyxy, track_id or None)
        if idx % a.every == 0:
            r = model.predict(frame, conf=a.conf, imgsz=a.imgsz, device=a.device, verbose=False)[0]
            det = r.boxes.cpu().numpy()
            n_det += 1
            tracked = tracker.update(det, frame)
            items += [(b[:4], int(b[4])) for b in tracked]  # confirmed tracks
            items += [(b, None) for b in det.xyxy]  # raw detections: don't wait for track confirmation
            items += [(t.xyxy, t.track_id) for t in tracker.tracked_stracks if not t.is_activated]
            items += [(t.xyxy, t.track_id) for t in tracker.lost_stracks if tracker.frame_id - t.end_frame <= a.hold]
            if keeper:  # recognise on the clean frame, before anything is drawn
                keeper.update(frame, idx, [(b[:4], int(b[4])) for b in tracked])
        else:
            # no detector this frame: advance Kalman state of every live track and use the prediction
            tracker.frame_id += 1
            pool = tracker.tracked_stracks + tracker.lost_stracks
            if pool:
                tracker.multi_predict(pool)
            items += [(t.xyxy, t.track_id) for t in tracker.tracked_stracks]
            items += [(t.xyxy, t.track_id) for t in tracker.lost_stracks if tracker.frame_id - t.end_frame <= a.hold]

        if keeper:
            kept_boxes = [np.asarray(b, dtype=float)[:4] for b, tid in items
                          if tid is not None and keeper.is_kept(tid, idx)]
            items = [(b, tid) for b, tid in items
                     if not (tid is not None and keeper.is_kept(tid, idx))
                     and not (tid is None and any(iou(np.asarray(b, dtype=float)[:4], kb) >= 0.5 for kb in kept_boxes))]

        for b, _ in items:
            box = pad_box(*np.asarray(b, dtype=float)[:4], a.pad, W, H)
            if a.mode == "black":
                blackout(frame, box)
            else:
                pixelate(frame, box, a.blocks)
        wq.put(frame)

        idx += 1
        if idx % 100 == 0:
            el = time.perf_counter() - t0
            print(f"\r{idx}/{total or '?'} frames  {idx / el:.1f} FPS", end="", flush=True)

    wq.put(None)
    wthread.join()
    cap.release()
    try:
        writer.stdin.close()
    except (BrokenPipeError, OSError):
        pass
    rc = writer.wait()
    if write_failed:
        rc = rc or 1
    el = time.perf_counter() - t0
    print(f"\rdone: {idx} frames in {el:.1f}s = {idx / el:.1f} FPS processing "
          f"(detector ran on {n_det} frames, every {a.every}) -> {a.output}")
    if keeper:
        print(f"keep: {len(keeper.ever_kept)} of {len(keeper.seen)} tracks were recognised as the target and left visible")
    if rc != 0:
        sys.exit(f"ffmpeg exited with code {rc}")


if __name__ == "__main__":
    main()