add video progress bar and save annotated video button

.gitignore

@@ -126,6 +126,7 @@ models/*
 
 # All cached movie files
 *.mp4
+*.MOV
 
 # All cached model files
 *.pt

app.py

@@ -2,6 +2,8 @@ from ultralytics import YOLO
 import time
 import os
 import logging
+import tempfile
+
 
 import av
 import cv2
@@ -138,7 +140,7 @@ def analyze_frame(frame: np.ndarray):
 
     if show_labels in ["Object Detection", "Both"]:
         # Run YOLOv8 object detection on the frame
-        object_results = object_model(frame)
+        object_results = object_model(frame, conf=0.5)
 
         for i, box in enumerate(object_results[0].boxes):
             class_id = int(box.cls)
@@ -151,7 +153,7 @@ def analyze_frame(frame: np.ndarray):
 
     if show_labels in ["Pose Estimation", "Both"]:
         # Run YOLOv8 pose estimation on the frame
-        pose_results = pose_model(frame)
+        pose_results = pose_model(frame, conf=0.5)
 
         for i, box in enumerate(pose_results[0].boxes):
             class_id = int(box.cls)
@@ -381,7 +383,7 @@ st.markdown(
 
 
 def analysis_init():
-    global yt_error, analysis_time, show_labels, labels_placeholder, input_subheader, input_placeholder, output_placeholder
+    global progress_bar, status_text, download_button, yt_error, analysis_time, show_labels, labels_placeholder, input_subheader, input_placeholder, output_placeholder
 
     yt_error.empty()  # Placeholder for analysis time
 
@@ -397,7 +399,11 @@ def analysis_init():
             ("Object Detection", "Pose Estimation", "Both"),
             index=2  # Set default to "Both" (index 2)
         )
+        # Create a progress bar
+        progress_bar = st.empty()
+        status_text = st.empty()
         labels_placeholder = st.empty()  # Placeholder for labels
+        download_button = st.empty()  # Placeholder for download button
 
 
 # Function to publish frames and results to the Streamlit UI
@@ -460,8 +466,25 @@ if uploaded_file is not None or image_url:
 # Function to process video files
 # This function reads frames from a video file, analyzes each frame for face detection and sentiment analysis,
 # and updates the Streamlit UI with the current input frame, analyzed frame, and detected labels.
+# Function to process video files
 def process_video(video_path):
     cap = cv2.VideoCapture(video_path)  # Open the video file
+
+    # Create a temporary file for the annotated video
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') as temp_video:
+        temp_video_path = temp_video.name
+
+    # save_annotated_video(video_path, temp_video_path)
+
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(temp_video_path, fourcc, fps, (width, height))
+
+    frame_count = 0
     while cap.isOpened():
         ret, frame = cap.read()  # Read a frame from the video
         if not ret:
@@ -472,35 +495,34 @@ def process_video(video_path):
 
         # Analyze the frame for face detection and sentiment analysis
         analyze_frame(rgb_frame)
-        publish_frame()  # Publish the results
+        analyzed_frame = img_container["analyzed"]
 
-    cap.release()  # Release the video capture object
+        if analyzed_frame is not None:
+            out.write(cv2.cvtColor(analyzed_frame, cv2.COLOR_RGB2BGR))
 
+        publish_frame()  # Publish the results
 
-# Function to get the video stream URL from YouTube using yt-dlp
+        # Update progress
+        frame_count += 1
+        progress = min(100, int(frame_count / total_frames * 100))
+        progress_bar.progress(progress)
+        status_text.text(f"Processing video: {progress}% complete")
 
-def get_youtube_stream_url(youtube_url):
-    ydl_opts = {
-        'format': 'bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best',
-        'quiet': True,
-        'no_warnings': True,
-    }
+    cap.release()  # Release the video capture object
+    out.release()
+
+    # Add download button for annotated video
+    with open(temp_video_path, "rb") as file:
+        download_button.download_button(
+            label="Download Annotated Video",
+            data=file,
+            file_name="annotated_video.mp4",
+            mime="video/mp4"
+        )
+
+    # Clean up the temporary file
+    os.unlink(temp_video_path)
 
-    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
-        try:
-            info_dict = ydl.extract_info(youtube_url, download=False)
-            if 'url' in info_dict:
-                return info_dict['url']
-            elif 'entries' in info_dict:
-                return info_dict['entries'][0]['url']
-            else:
-                yt_error.error(
-                    "Unable to extract video URL. The video might be unavailable or restricted.")
-                return None
-        except yt_dlp.utils.DownloadError as e:
-            yt_error.error(
-                f"Error: Unable to process the YouTube URL. {str(e)}")
-            return None
 
 # Function to get video URL using Cobalt API
 

requirements.txt

@@ -6,4 +6,5 @@ setuptools
 streamlit_webrtc
 twilio
 ultralytics
-watchdog
+watchdog
+tempfile