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