Added YouTube URL option

app.py

@@ -1,11 +1,8 @@
+import yt_dlp
 from ultralytics import YOLO
-import torch
-import tensorflow as tf
 import time
 import os
 import logging
-from pathlib import Path
-from typing import List
 
 import av
 import cv2
@@ -16,7 +13,6 @@ from streamlit_webrtc import WebRtcMode, webrtc_streamer
 from utils.download import download_file
 from utils.turn import get_ice_servers
 
-from mtcnn import MTCNN  # Import MTCNN for face detection
 from PIL import Image, ImageDraw  # Import PIL for image processing
 from transformers import pipeline  # Import Hugging Face transformers pipeline
 
@@ -98,17 +94,6 @@ def analyze_frame(frame: np.ndarray):
 # Suppress FFmpeg logs
 os.environ["FFMPEG_LOG_LEVEL"] = "quiet"
 
-# Suppress TensorFlow or PyTorch progress bars
-
-tf.get_logger().setLevel("ERROR")
-os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-
-# Suppress PyTorch logs
-
-logging.getLogger().setLevel(logging.WARNING)
-torch.set_num_threads(1)
-logging.getLogger("torch").setLevel(logging.ERROR)
-
 # Suppress Streamlit logs using the logging module
 logging.getLogger("streamlit").setLevel(logging.ERROR)
 
@@ -205,6 +190,10 @@ with col1:
     st.subheader("Or Enter Image URL")
     image_url = st.text_input("Image URL")
 
+    # Text input for video URL
+    st.subheader("Enter a YouTube URL")
+    youtube_url = st.text_input("YouTube URL")
+
     # File uploader for videos
     st.subheader("Upload a Video")
     uploaded_video = st.file_uploader(
@@ -324,6 +313,28 @@ def process_video(video_path):
     cap.release()  # Release the video capture object
 
 
+# Function to get the video stream URL from YouTube using yt-dlp
+
+
+def get_youtube_stream_url(youtube_url):
+    ydl_opts = {
+        'format': 'best[ext=mp4]',
+        'quiet': True,
+    }
+    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+        info_dict = ydl.extract_info(youtube_url, download=False)
+        stream_url = info_dict['url']
+    return stream_url
+
+
+# If a YouTube URL is provided, process the video
+if youtube_url:
+    analysis_init()  # Initialize the analysis UI
+
+    stream_url = get_youtube_stream_url(youtube_url)
+
+    process_video(stream_url)  # Process the video
+
 # If a video is uploaded or a URL is provided, process the video
 if uploaded_video is not None or video_url:
     analysis_init()  # Initialize the analysis UI

requirements.txt

@@ -2,10 +2,9 @@ streamlit
 opencv-python-headless
 numpy
 transformers
-torch
 setuptools
-tensorflow
-tf-keras
 streamlit_webrtc
 twilio
-ultralytics
+ultralytics
+yt-dlp
+watchdog