Update app.py

app.py

@@ -19,7 +19,6 @@ Dependencies:
 Usage:
 Run this file to launch the Gradio interface, which allows users to input search queries for YouTube live streams, select a stream, and perform object detection on the selected live stream.
 """
-
 import logging
 import sys
 from enum import Enum
@@ -31,11 +30,10 @@ import innertube
 import numpy as np
 from PIL import Image
 from ultralytics import YOLO
-from pytube import YouTube
+import yt_dlp
 
 logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
 
-
 class SearchFilter(Enum):
     LIVE = ("EgJAAQ%3D%3D", "Live")
     VIDEO = ("EgIQAQ%3D%3D", "Video")
@@ -47,7 +45,6 @@ class SearchFilter(Enum):
     def __str__(self):
         return self.human_readable
 
-
 class SearchService:
     @staticmethod
     def search(query: Optional[str], filter: SearchFilter = SearchFilter.VIDEO):
@@ -93,37 +90,40 @@ class SearchService:
 
     @staticmethod
     def get_stream(youtube_url: str) -> Optional[str]:
-        """Retrieves the live stream URL for a given YouTube video URL using pytube.
-    
+        """Retrieves the livestream URL for a given YouTube video URL using yt-dlp.
+
         :param youtube_url: The URL of the YouTube video.
         :type youtube_url: str
-        :return: The HLS manifest URL if available, otherwise None.
+        :return: The livestream URL if available, otherwise None.
         :rtype: Optional[str]
         """
+        ydl_opts = {
+            'format': 'best',
+            'quiet': True,
+            'no_warnings': True,
+            'force_generic_extractor': False,
+            'skip_download': True,
+        }
         try:
-            yt = YouTube(youtube_url)
-            _ = yt.streams
-            is_live = yt.player_response.get('videoDetails', {}).get('isLiveContent', False)
-            if 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
+            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                info_dict = ydl.extract_info(youtube_url, download=False)
+                if info_dict.get('is_live'):
+                    live_url = info_dict.get('url')
+                    if live_url:
+                        logging.debug(f"Found livestream URL: {live_url}")
+                        return live_url
+                    else:
+                        logging.warning(f"Livestream URL not found for: {youtube_url}")
+                        return None
                 else:
-                    logging.warning(f"HLS manifest URL not found for live stream: {youtube_url}")
+                    logging.warning(f"Video is not a livestream: {youtube_url}")
                     return None
-            else:
-                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}")
             return None
 
-
 INITIAL_STREAMS = SearchService.search("world live cams", SearchFilter.LIVE)
 
-
 class LiveYouTubeObjectDetector:
     def __init__(self):
         logging.getLogger().setLevel(logging.DEBUG)
@@ -163,36 +163,36 @@ class LiveYouTubeObjectDetector:
             return None
         try:
             cap = cv2.VideoCapture(stream_url)
+            if not cap.isOpened():
+                logging.warning("Failed to open video capture.")
+                return None
             ret, frame = cap.read()
             cap.release()
             if ret and frame is not None:
-                return cv2.resize(frame, (1920, 1080))
+                return cv2.resize(frame, (1280, 720))
             else:
-                logging.warning("Unable to process the live stream with cv2.VideoCapture.")
+                logging.warning("Unable to read frame from the live stream.")
                 return None
         except Exception as e:
             logging.warning(f"An error occurred while capturing the frame: {e}")
             return None
 
     def annotate(self, frame: np.ndarray) -> Tuple[Image.Image, List[Tuple[Tuple[int, int, int, int], str]]]:
-        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        predictions = self.model.predict(frame_rgb)
+        results = self.model(frame)[0]
         annotations = []
-        result = predictions[0]
-
-        boxes = result.boxes
+        boxes = results.boxes
         for box in boxes:
             x1, y1, x2, y2 = box.xyxy[0].tolist()
             class_id = int(box.cls[0])
             class_name = self.model.names[class_id]
             bbox_coords = (int(x1), int(y1), int(x2), int(y2))
             annotations.append((bbox_coords, class_name))
-
-        return Image.fromarray(frame_rgb), annotations
+        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+        return pil_image, annotations
 
     @staticmethod
     def create_black_image() -> Tuple[Image.Image, List]:
-        black_image = np.zeros((1080, 1920, 3), dtype=np.uint8)
+        black_image = np.zeros((720, 1280, 3), dtype=np.uint8)
         pil_black_image = Image.fromarray(black_image)
         return pil_black_image, []
 
@@ -240,6 +240,5 @@ class LiveYouTubeObjectDetector:
 
         app.queue().launch(show_api=False, debug=True)
 
-
 if __name__ == "__main__":
     LiveYouTubeObjectDetector().render()