File size: 2,353 Bytes
c814e57 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import streamlit as st
import cv2
import tempfile
from ultralytics import YOLO
import numpy as np
alerting_classes = {
0: 'People',
2: 'Car',
7: 'Truck',
24: 'Backpack',
65: 'Suspicious handheld device',
26: 'Handbag',
28: 'Suitcase',
}
red_tint = np.array([[[0, 0, 255]]], dtype=np.uint8)
model1 = YOLO('yolov8n.pt')
st.title("Object Detection and Recognition")
video_file = st.file_uploader("Choose a video file", type=["mp4"])
if video_file is not None:
# Create temporary file for uploaded video
tfile = tempfile.NamedTemporaryFile(delete=False)
tfile.write(video_file.read())
# Open video capture using temporary file path
cap = cv2.VideoCapture(tfile.name)
alert_set = set(alerting_classes.keys())
alert_set.remove(0)
# Create red-tinted overlay
red_tinted_overlay = np.tile(red_tint, (1, 1, 1))
stframe = st.empty()
stop_button = st.button("Stop Inference")
while cap.isOpened() and not stop_button:
alert_flag = False
alert_reason = []
success, frame = cap.read()
# if frame is read correctly ret is True
if not success:
st.warning("Can't receive frame (stream end?). Exiting ...")
break
if success:
# Perform YOLO object detection
results = model1(frame, conf=0.35, verbose=False, classes=list(alerting_classes.keys()))
class_ids = results[0].boxes.cls.tolist()
class_counts = {cls: class_ids.count(cls) for cls in set(class_ids)}
for cls in alert_set:
if cls in class_counts and class_counts[cls] > 0:
alert_flag = True
alert_reason.append((cls, class_counts[cls]))
if class_counts.get(0, 0) > 5:
alert_flag = True
alert_reason.append((0, class_counts[0]))
# Draw bounding boxes and alerts if necessary
img = results[0].plot()
if alert_flag:
red_tinted_overlay = cv2.resize(red_tinted_overlay, (img.shape[1], img.shape[0]))
img = cv2.addWeighted(img, 0.7, red_tinted_overlay, 0.3, 0)
stframe.image(img, channels="BGR", caption="YOLOv8 Inference")
del results
cap.release()
cv2.destroyAllWindows()
tfile.close()
|