82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
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 prometheus_client import Gauge, start_http_server
|
||
|
||
# Nur einmal beim ersten Start starten
|
||
if "prom" not in st.session_state:
|
||
# Starte Prometheus-Metrik-Server (z. B. auf Port 9100)
|
||
start_http_server(9100)
|
||
|
||
# Metriken definieren
|
||
st.session_state.prom = {
|
||
'people_in': Gauge("people_in_count", "Number of people who entered"),
|
||
'people_out': Gauge("people_out_count", "Number of people who exited"),
|
||
}
|
||
|
||
st.set_page_config(layout="wide", page_title="People Counter")
|
||
|
||
# RTSP-Stream URL
|
||
#RTSP_URL = st.text_input("RTSP Stream URL", "rtsp://<benutzer>:<passwort>@<ip>:<port>/pfad")
|
||
RTSP_URL = "rtsp://localhost:8554/cam"
|
||
|
||
start_button = st.button("Start Counter")
|
||
|
||
FRAME_SKIP = 2 # Reduziert Verarbeitungslast
|
||
|
||
if start_button and RTSP_URL:
|
||
line_orientation = 'vertical'
|
||
line_position = 640
|
||
|
||
stframe = st.empty()
|
||
model = YOLO("yolo_weights/yolo11n.pt")
|
||
counter = Counter(line_orientation)
|
||
|
||
cap = cv2.VideoCapture(RTSP_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)
|
||
st.session_state.prom['people_in'].set(counter.in_count)
|
||
st.session_state.prom['people_out'].set(counter.out_count)
|
||
|
||
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()
|