69 lines
2.3 KiB
Python
69 lines
2.3 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 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()
|