From 91f4d70a77cb8ea3ef5388498903f5e9922b9b97 Mon Sep 17 00:00:00 2001 From: Jakob Lechner Date: Sat, 2 Aug 2025 03:25:54 +0200 Subject: [PATCH] Replace streamlit with FastAPI --- main.py | 75 +++++++++++++++++++++++++++ metrics.py | 15 ++++-- requirements.txt | 3 +- streamlit_app.py | 69 ------------------------- tracker.py | 130 +++++++++++++++++++++++++++++++++++++++++++++++ utils/counter.py | 16 +++--- utils/zones.py | 2 + 7 files changed, 228 insertions(+), 82 deletions(-) create mode 100644 main.py delete mode 100644 streamlit_app.py create mode 100644 tracker.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..e0b3d53 --- /dev/null +++ b/main.py @@ -0,0 +1,75 @@ +from fastapi import FastAPI, Request +from fastapi.responses import HTMLResponse, StreamingResponse, PlainTextResponse +from tracker import start_all_streams, get_metrics, get_latest_frame, STREAMS +from prometheus_client import generate_latest, CONTENT_TYPE_LATEST +import time + +app = FastAPI() + + +@app.on_event("startup") +def startup_event(): + start_all_streams() + + +@app.get("/", response_class=HTMLResponse) +def index(request: Request): + current_stream = request.query_params.get("stream") + + stream_options = "\n".join( + [ + f'' + for stream_id in STREAMS.keys() + ] + ) + + return f""" + + +

People Counter Dashboard

+
+ + +
+ +

Prometheus Metrics

+ + + """ + + +@app.get("/metrics") +def metrics(): + return PlainTextResponse( + generate_latest(get_metrics()), media_type=CONTENT_TYPE_LATEST + ) + + +@app.get("/video_feed/{stream_id}") +def video_feed(stream_id: str): + boundary = "--boundarydonotcross" + + def generate(): + last_version = time.time() + while True: + frame, version = get_latest_frame(stream_id) + + if frame is not None and version > last_version: + last_version = version + print(f"yielding new frame @{version}") + yield ( + boundary.encode() + + b"\r\n" + + b"Content-Type: image/jpeg\r\n\r\n" + + frame + + b"\r\n" + ) + else: + # wait for new frame + time.sleep(0.05) + + return StreamingResponse( + generate(), media_type=f"multipart/x-mixed-replace; boundary={boundary}" + ) diff --git a/metrics.py b/metrics.py index 4be4b2f..0bed927 100644 --- a/metrics.py +++ b/metrics.py @@ -1,23 +1,28 @@ import threading -from prometheus_client import Counter, start_http_server +from prometheus_client import Counter, CollectorRegistry _lock = threading.Lock() _registry = None _metrics = None + def get_metrics(): - global _metrics + global _registry, _metrics with _lock: - if _metrics is None: - start_http_server(9110) + if _registry is None: + _registry = CollectorRegistry() _metrics = { "people_in": Counter( "people_in_count", "Number of people who entered", + ["stream"], + registry=_registry, ), "people_out": Counter( "people_out_count", "Number of people who exited", + ["stream"], + registry=_registry, ), } - return _metrics + return _registry, _metrics diff --git a/requirements.txt b/requirements.txt index 6826106..abf93cc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,4 +4,5 @@ numpy filterpy scikit-image supervision -streamlit +uvicorn +fastapi diff --git a/streamlit_app.py b/streamlit_app.py deleted file mode 100644 index f2afa7f..0000000 --- a/streamlit_app.py +++ /dev/null @@ -1,69 +0,0 @@ -import streamlit as st -import cv2 -from ultralytics import YOLO -from utils.counter import Counter -from utils.zones import draw_count_line_horizontal, draw_count_line_vertical -import tempfile -import time - -from metrics import get_metrics - -metrics = get_metrics() - -st.set_page_config(layout="wide", page_title="People Counter") - -STREAM_URL = "http://192.168.11.76:8080?action=stream" - -start_button = st.button("Start Counter") - -FRAME_SKIP = 2 # Reduziert Verarbeitungslast - -if start_button and STREAM_URL: - line_orientation = 'vertical' - line_position = 640 - - stframe = st.empty() - model = YOLO("yolo_weights/yolo11n.pt") - counter = Counter(line_orientation, metrics) - - cap = cv2.VideoCapture(STREAM_URL) - - frame_idx = 0 - while cap.isOpened(): - ret, frame = cap.read() - if not ret: - st.warning("Kein Bild vom RTSP-Stream.") - break - - frame_idx += 1 - if frame_idx % FRAME_SKIP != 0: - continue - - results = model.track(frame, persist=True, classes=0, tracker="bytetrack.yaml") - boxes = results[0].boxes - - if boxes.id is not None: - for box, track_id in zip(boxes.xywh.cpu(), boxes.id.cpu()): - x, y, w, h = map(int, box) - cv2.rectangle(frame, (x-w//2, y-h//2), (x+w//2, y+h//2), (0, 255, 0), 2) - cv2.putText(frame, str(track_id), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0,255,0), 2) - - tracks = [ - type("Track", (), {"id": int(track_id), "bbox": box.numpy()}) - for box, track_id in zip(boxes.xywh.cpu(), boxes.id.cpu()) - ] - counter.update(tracks, line_position) - - if line_orientation == 'horizontal': - draw_count_line_horizontal(frame, line_position) - elif line_orientation == 'vertical': - draw_count_line_vertical(frame, line_position) - else: - raise NotImplementedError(f'Line orientation {line_orientation} is invalid!') - cv2.putText(frame, f"In: {counter.in_count}", (10, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2) - cv2.putText(frame, f"Out: {counter.out_count}", (10, 80), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 255), 2) - - frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - stframe.image(frame, channels="RGB") - - cap.release() diff --git a/tracker.py b/tracker.py new file mode 100644 index 0000000..b11ca20 --- /dev/null +++ b/tracker.py @@ -0,0 +1,130 @@ +import cv2 +import threading +import time +from ultralytics import YOLO +from utils.counter import Counter +from utils.zones import draw_count_line_vertical +from collections import defaultdict + +from metrics import get_metrics + +registry, metrics = get_metrics() + +STREAMS = { + "Kasse 1": "http://192.168.11.76:8080?action=stream", + "Kasse 2": "http://192.168.11.230:8080?action=stream", +} + +FRAME_SKIP = 2 +model = YOLO("yolo_weights/yolo11n.pt") + +latest_frames = {} + + +def process_stream(stream_id, url): + print(f"PROCESS STREAM {stream_id} from {url}") + line_orientation = "vertical" + counter = Counter(stream_id, line_orientation, metrics) + cap = cv2.VideoCapture(url) + frame_idx = 0 + + while cap.isOpened(): + ret, frame = cap.read() + if not ret: + time.sleep(1) + continue + + height, width, _ = frame.shape + + frame_idx += 1 + if frame_idx % FRAME_SKIP != 0: + continue + + try: + results = model.track( + frame, persist=True, classes=0, tracker="bytetrack.yaml" + ) + except Exception as e: + print(e) + continue + + if line_orientation == "horizontal": + line_position = int(height / 2) + elif line_orientation == "vertical": + line_position = int(width / 2) + else: + raise NotImplementedError( + f"Line orientation {line_orientation} is invalid!" + ) + + boxes = results[0].boxes + if boxes.id is not None: + for box, track_id in zip(boxes.xywh.cpu(), boxes.id.cpu()): + x, y, w, h = map(int, box) + cv2.rectangle( + frame, + (x - w // 2, y - h // 2), + (x + w // 2, y + h // 2), + (0, 255, 0), + 2, + ) + cv2.putText( + frame, + str(track_id), + (x, y - 10), + cv2.FONT_HERSHEY_SIMPLEX, + 0.6, + (0, 255, 0), + 2, + ) + + tracks = [ + type("Track", (), {"id": int(track_id), "bbox": box.numpy()}) + for box, track_id in zip(boxes.xywh.cpu(), boxes.id.cpu()) + ] + counter.update(tracks, line_position) + + if line_orientation == "horizontal": + draw_count_line_horizontal(frame, line_position) + elif line_orientation == "vertical": + draw_count_line_vertical(frame, line_position) + + cv2.putText( + frame, + f"In: {counter.in_count}", + (10, 40), + cv2.FONT_HERSHEY_SIMPLEX, + 1, + (0, 255, 255), + 2, + ) + cv2.putText( + frame, + f"Out: {counter.out_count}", + (10, 80), + cv2.FONT_HERSHEY_SIMPLEX, + 1, + (0, 255, 255), + 2, + ) + + _, jpeg = cv2.imencode(".jpg", frame) + latest_frames[stream_id] = (jpeg.tobytes(), time.time()) + print(f"[{stream_id}] Frame captured and stored") + + cap.release() + + +def start_all_streams(): + for stream_id, url in STREAMS.items(): + threading.Thread( + target=process_stream, args=(stream_id, url), daemon=True + ).start() + + +def get_metrics(): + return registry + + +def get_latest_frame(stream_id): + return latest_frames.get(stream_id, (None, None)) diff --git a/utils/counter.py b/utils/counter.py index 39b56e9..cf91a45 100644 --- a/utils/counter.py +++ b/utils/counter.py @@ -1,23 +1,25 @@ class Counter: - def __init__(self, line_orientation, metrics): + def __init__(self, stream_id, line_orientation, metrics): self.in_count = 0 self.out_count = 0 self.track_memory = {} self.line_orientation = line_orientation self.metrics = metrics + self.stream_id = stream_id def update(self, tracks, line_position): for track in tracks: track_id = track.id x, y, w, h = track.bbox - - if self.line_orientation == 'horizontal': + if self.line_orientation == "horizontal": center = int(y + h / 2) - elif self.line_orientation == 'vertical': + elif self.line_orientation == "vertical": center = int(x + w / 2) else: - raise NotImplementedError(f'Line orientation {self.line_orientation} is invalid!') + raise NotImplementedError( + f"Line orientation {self.line_orientation} is invalid!" + ) if track_id not in self.track_memory: self.track_memory[track_id] = center @@ -28,7 +30,7 @@ class Counter: if prev < line_position <= center: self.in_count += 1 - self.metrics["people_in"].inc() + self.metrics["people_in"].labels(stream=self.stream_id).inc() elif prev > line_position >= center: self.out_count += 1 - self.metrics["people_out"].inc() + self.metrics["people_out"].labels(stream=self.stream_id).inc() diff --git a/utils/zones.py b/utils/zones.py index 4878b21..c3efe3f 100644 --- a/utils/zones.py +++ b/utils/zones.py @@ -1,10 +1,12 @@ import cv2 + def draw_count_line_horizontal(frame, y): color = (0, 255, 255) thickness = 2 cv2.line(frame, (0, y), (frame.shape[1], y), color, thickness) + def draw_count_line_vertical(frame, x): color = (0, 255, 255) thickness = 2