Update app.py

app.py

@@ -1,5 +1,5 @@
 """
-This module integrates real-time object detection into live YouTube streams using the YOLO model, and provides an interactive user interface through Gradio. It allows users to search for live YouTube streams and apply object detection to these streams in real time.
+This module integrates real-time object detection into live YouTube streams using the YOLO model and provides an interactive user interface through Gradio. It allows users to search for live YouTube streams and apply object detection to these streams in real time.
 
 Main Features:
 - Search for live YouTube streams using specific queries.
@@ -31,8 +31,11 @@ import innertube
 import numpy as np
 from PIL import Image
 from ultralytics import YOLO
+
+# Import pytube
 from pytube import YouTube
 
+# Set up logging
 logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
 
 
@@ -93,22 +96,26 @@ class SearchService:
 
     @staticmethod
     def get_stream(youtube_url: str) -> Optional[str]:
-        """Retrieves the stream URL for a given YouTube video URL using pytube.
+        """Retrieves the live stream URL for a given YouTube video URL using pytube.
 
         :param youtube_url: The URL of the YouTube video.
         :type youtube_url: str
-        :return: The stream URL if available, otherwise None.
+        :return: The HLS manifest URL if available, otherwise None.
         :rtype: Optional[str]
         """
         try:
             yt = YouTube(youtube_url)
-            # Live streams have a 'LIVE' attribute in the 'streaming_data' key
-            if 'hlsManifestUrl' in yt.watch_html:
-                # Extract the HLS manifest URL
-                hls_manifest_url = yt.vid_info['streamingData']['hlsManifestUrl']
-                return hls_manifest_url
+            if yt.is_live:
+                streaming_data = yt.player_response.get('streamingData', {})
+                hls_manifest_url = streaming_data.get('hlsManifestUrl')
+                if hls_manifest_url:
+                    logging.debug(f"Found HLS manifest URL for live stream: {hls_manifest_url}")
+                    return hls_manifest_url
+                else:
+                    logging.warning(f"HLS manifest URL not found for live stream: {youtube_url}")
+                    return None
             else:
-                logging.warning(f"No live stream found for: {youtube_url}")
+                logging.warning(f"Video is not a live stream: {youtube_url}")
                 return None
         except Exception as e:
             logging.warning(f"An error occurred while getting stream: {e}")
@@ -156,7 +163,6 @@ class LiveYouTubeObjectDetector:
         if not stream_url:
             return None
         try:
-            # Use cv2 VideoCapture with the HLS manifest URL
             cap = cv2.VideoCapture(stream_url)
             ret, frame = cap.read()
             cap.release()
@@ -194,10 +200,11 @@ class LiveYouTubeObjectDetector:
     @staticmethod
     def get_live_streams(query=""):
         return SearchService.search(query if query else "world live cams", SearchFilter.LIVE)
-    
+
     def render(self):
         with gr.Blocks(title="Object Detection in Live YouTube Streams",
-                       css="footer {visibility: hidden}", analytics_enabled=False) as app:
+                       css="footer {visibility: hidden}",
+                       analytics_enabled=False) as app:
             self.page_title.render()
             with gr.Column():
                 with gr.Group():
@@ -211,7 +218,7 @@ class LiveYouTubeObjectDetector:
                     self.search_button.render()
             with gr.Row():
                 self.gallery.render()
-    
+
             @self.gallery.select(inputs=None, outputs=[self.annotated_image, self.stream_input], scroll_to_output=True)
             def detect_objects_from_gallery_item(evt: gr.SelectData):
                 if evt.index is not None and evt.index < len(self.streams):
@@ -221,18 +228,19 @@ class LiveYouTubeObjectDetector:
                     annotated_image_result = self.detect_objects(stream_url)
                     return annotated_image_result, stream_url
                 return self.create_black_image(), ""
-    
+
             @self.search_button.click(inputs=[self.search_input], outputs=[self.gallery])
             def search_live_streams(query):
                 self.streams = self.get_live_streams(query)
                 gallery_items = [(stream["thumbnail_url"], stream["title"]) for stream in self.streams]
                 return gallery_items
-    
+
             @self.submit_button.click(inputs=[self.stream_input], outputs=[self.annotated_image])
             def detect_objects_from_url(url):
                 return self.detect_objects(url)
-    
+
         app.queue().launch(show_api=False, debug=True)
 
+
 if __name__ == "__main__":
     LiveYouTubeObjectDetector().render()