Create app.py

app.py

@@ -0,0 +1,80 @@
+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()