Update app.py

app.py

@@ -3,6 +3,7 @@ import cv2
 import tempfile
 from ultralytics import YOLO
 import numpy as np
+import time
 
 alerting_classes = {
     0: 'People',
@@ -35,22 +36,32 @@ if video_file is not None:
     # Create red-tinted overlay
     red_tinted_overlay = np.tile(red_tint, (1, 1, 1))
 
-    stframe = st.empty()
-
     stop_button = st.button("Stop Inference")
+    processing_interrupted = False
+
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    progress_bar_processing_slot = st.empty()
+
+    # Collect frames in a list
+    frames = []
 
-    while cap.isOpened() and not stop_button:
+    while cap.isOpened() and not processing_interrupted:
         alert_flag = False
         alert_reason = []
 
         success, frame = cap.read()
 
-        # if frame is read correctly ret is True
+        # if the frame is read correctly ret is True
         if not success:
-            st.warning("Can't receive frame (stream end?). Exiting ...")
+            # st.warning("Can't receive frame (stream end?). Exiting ...")
             break
 
         if success:
+            # Check if the stop button is clicked
+            if stop_button:
+                processing_interrupted = True
+                break
+
             # Perform YOLO object detection
             results = model1(frame, conf=0.35, verbose=False, classes=list(alerting_classes.keys()))
 
@@ -66,15 +77,50 @@ if video_file is not None:
                 alert_flag = True
                 alert_reason.append((0, class_counts[0]))
 
-            # Draw bounding boxes and alerts if necessary
+            text = 'ALERT!'
+            font = cv2.FONT_HERSHEY_DUPLEX
+            font_scale = 0.75
+            thickness = 2
+
+            size = cv2.getTextSize(text, font, font_scale, thickness)
+            x = 0
+            y = int((2 + size[0][1]))
+
             img = results[0].plot()
             if alert_flag:
+                # Resize the red-tinted overlay to match the image size
                 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)
+                cv2.putText(img, text, (x, y), font, font_scale, (0, 0, 0), thickness)
 
-            stframe.image(img, channels="BGR", caption="YOLOv8 Inference")
+                y += int(size[0][1]) + 10  # Move to the next line
+
+                for cls, count in alert_reason:
+                    alert_text = f'{count} {alerting_classes[cls]}'
+                    cv2.putText(img, alert_text, (x, y), font, font_scale, (0, 0, 0), thickness)
+                    y += int(size[0][1]) + 10  # Move to the next line
+
+            # Append the frame to the list
+            frames.append(img)
 
             del results
+
+            # Update processing progress bar
+            current_frame_processing = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
+            progress = current_frame_processing / total_frames
+            progress_bar_processing_slot.progress(progress)
+            progress_bar_processing_slot.text(f"Processing... {int(progress * 100)}%")
+
+    # Release resources
     cap.release()
-    cv2.destroyAllWindows()
     tfile.close()
+
+    # Display frames one by one as a video with 24 FPS
+    if processing_interrupted:
+        st.text("User interrupted processing.")
+    else:
+        progress_bar_processing_slot.text("Done!")
+        video_placeholder = st.image([], channels="BGR", caption="YOLOv8 Inference")
+        progress_bar_display = st.progress(0)
+        fps_delay = 1 / 24  # Delay to achieve 24 FPS
+