● Zion Boggan
repos/Sightline/backend/app/services/video_alpr.py
zionboggan.com ↗
153 lines · python
History for this file →
1
"""Video / RTSP frame path for ALPR. Samples frames from a video file or an RTSP/HTTP stream
2
(both via OpenCV's VideoCapture), runs the plate recognizer on each sampled frame, and aggregates
3
detections so each distinct plate becomes ONE read (best OCR across the frames it appeared in) -
4
not one read per frame. This is the temporal layer on top of the single-image engine.
5
"""
6
from __future__ import annotations
7
 
8
import os
9
import tempfile
10
import time
11
from dataclasses import dataclass, field
12
 
13
from app.services.alpr import recognize_array
14
 
15
 
16
@dataclass
17
class PlateTrack:
18
    plate_norm: str
19
    plate: str
20
    best_ocr: float
21
    best_det: float
22
    frames_seen: int
23
    first_ts: float
24
    last_ts: float
25
    bbox: list[int] = field(default_factory=list)
26
 
27
 
28
def _aggregate(per_frame: list[tuple[float, list]]) -> list[PlateTrack]:
29
    """per_frame: list of (timestamp_seconds, [PlateDetection,...]). Collapse by plate_norm."""
30
    tracks: dict[str, PlateTrack] = {}
31
    for ts, dets in per_frame:
32
        for d in dets:
33
            t = tracks.get(d.plate_norm)
34
            if t is None:
35
                tracks[d.plate_norm] = PlateTrack(
36
                    plate_norm=d.plate_norm, plate=d.plate, best_ocr=d.ocr_confidence,
37
                    best_det=d.detection_confidence, frames_seen=1, first_ts=ts, last_ts=ts,
38
                    bbox=d.bbox)
39
            else:
40
                t.frames_seen += 1
41
                t.last_ts = ts
42
                if d.ocr_confidence > t.best_ocr:
43
                    t.best_ocr, t.plate, t.bbox = d.ocr_confidence, d.plate, d.bbox
44
                t.best_det = max(t.best_det, d.detection_confidence)
45
    return sorted(tracks.values(), key=lambda t: (t.frames_seen, t.best_ocr), reverse=True)
46
 
47
 
48
def process_capture(cap, sample_fps: float = 2.0, max_frames: int = 240,
49
                    max_seconds: float | None = None, min_conf: float = 0.4) -> dict:
50
    """Drive an opened cv2.VideoCapture: sample ~sample_fps frames/sec, recognize, aggregate.
51
    Works identically for files and live streams."""
52
    import cv2
53
 
54
    src_fps = cap.get(cv2.CAP_PROP_FPS) or 0
55
    stride = max(1, int(round(src_fps / sample_fps))) if src_fps and src_fps > 0 else 0
56
    per_frame: list[tuple[float, list]] = []
57
    processed = sampled = 0
58
    start = time.monotonic()
59
    last_sample = 0.0
60
    idx = 0
61
    while True:
62
        ok, frame = cap.read()
63
        if not ok:
64
            break
65
        idx += 1
66
        now = time.monotonic() - start
67
        if max_seconds and now > max_seconds:
68
            break
69
        take = False
70
        if stride:
71
            take = (idx % stride == 0)
72
            ts = idx / src_fps
73
        else:
74
            if now - last_sample >= 1.0 / sample_fps:
75
                take, last_sample, ts = True, now, now
76
            else:
77
                ts = now
78
        if not take:
79
            continue
80
        sampled += 1
81
        try:
82
            dets = recognize_array(frame, min_conf=min_conf)
83
        except Exception:
84
            dets = []
85
        per_frame.append((round(ts, 2), dets))
86
        processed += 1
87
        if processed >= max_frames:
88
            break
89
 
90
    tracks = _aggregate(per_frame)
91
    return {"frames_sampled": sampled, "frames_with_detections": sum(1 for _, d in per_frame if d),
92
            "plates": [{"plate_norm": t.plate_norm, "plate": t.plate,
93
                        "ocr_confidence": round(t.best_ocr, 3), "detection_confidence": round(t.best_det, 3),
94
                        "frames_seen": t.frames_seen, "first_ts": t.first_ts, "last_ts": t.last_ts,
95
                        "bbox": t.bbox} for t in tracks]}
96
 
97
 
98
def process_video_bytes(data: bytes, sample_fps: float = 2.0, max_frames: int = 240,
99
                        min_conf: float = 0.4) -> dict:
100
    import cv2
101
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tf:
102
        tf.write(data)
103
        path = tf.name
104
    try:
105
        cap = cv2.VideoCapture(path)
106
        if not cap.isOpened():
107
            raise ValueError("could not open video (unsupported codec?)")
108
        try:
109
            return process_capture(cap, sample_fps=sample_fps, max_frames=max_frames, min_conf=min_conf)
110
        finally:
111
            cap.release()
112
    finally:
113
        try:
114
            os.remove(path)
115
        except OSError:
116
            pass
117
 
118
 
119
def open_stream(url: str):
120
    """Open an RTSP/HTTP(S)/file stream. Returns an opened VideoCapture or raises."""
121
    import cv2
122
    if url.lower().startswith("rtsp"):
123
        os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = "rtsp_transport;tcp"
124
    else:
125
        os.environ.pop("OPENCV_FFMPEG_CAPTURE_OPTIONS", None)
126
    cap = cv2.VideoCapture(url, cv2.CAP_FFMPEG)
127
    if not cap.isOpened():
128
        cap.release()
129
        raise ValueError("could not open stream (check URL / reachability / codec)")
130
    return cap
131
 
132
 
133
def probe_stream(url: str) -> dict:
134
    """Connect and grab a single frame to verify a stream is reachable and decodable."""
135
    cap = open_stream(url)
136
    try:
137
        ok, frame = cap.read()
138
        if not ok or frame is None:
139
            return {"reachable": True, "decoded": False}
140
        h, w = frame.shape[:2]
141
        return {"reachable": True, "decoded": True, "width": int(w), "height": int(h)}
142
    finally:
143
        cap.release()
144
 
145
 
146
def process_stream(url: str, seconds: float = 15.0, sample_fps: float = 2.0,
147
                   max_frames: int = 120, min_conf: float = 0.4) -> dict:
148
    cap = open_stream(url)
149
    try:
150
        return process_capture(cap, sample_fps=sample_fps, max_frames=max_frames,
151
                               max_seconds=seconds, min_conf=min_conf)
152
    finally:
153
        cap.release()