Update app.py

app.py

@@ -1,15 +1,14 @@
 import logging
-import os
 import sys
-import zipfile
 from enum import Enum
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Optional
+
 import cv2
 import gradio as gr
 import innertube
 import numpy as np
 import streamlink
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image
 from ultralytics import YOLO
 
 logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
@@ -28,40 +27,34 @@ class SearchFilter(Enum):
     def __str__(self):
         return self.human_readable
 
+
 class SearchService:
     @staticmethod
-    def search(
-        query: Optional[str], filter: SearchFilter = SearchFilter.VIDEO
-    ) -> (List[Dict[str, Any]], Optional[str]):
+    def search(query: Optional[str], filter: SearchFilter = SearchFilter.VIDEO):
         client = innertube.InnerTube("WEB", "2.20230920.00.00")
         response = SearchService._search(query, filter)
         results = SearchService.parse(response)
         return results
 
     @staticmethod
-    def parse(data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def parse(data: Dict[str, Any]):
         results = []
         contents = data["contents"]["twoColumnSearchResultsRenderer"]["primaryContents"]["sectionListRenderer"]["contents"]
         items = contents[0]["itemSectionRenderer"]["contents"] if contents else []
         for item in items:
             if "videoRenderer" in item:
                 renderer = item["videoRenderer"]
-                video_id = renderer["videoId"]
-                title = "".join(run["text"] for run in renderer["title"]["runs"])
-                thumbnail_url = renderer["thumbnail"]["thumbnails"][-1]["url"]
-                result = {
-                    "video_id": video_id,
-                    "thumbnail_url": thumbnail_url,
-                    "title": title,
-                }
-                results.append(result)
-
+                results.append(
+                    {
+                        "video_id": renderer["videoId"],
+                        "thumbnail_url": renderer["thumbnail"]["thumbnails"][-1]["url"],
+                        "title": "".join(run["text"] for run in renderer["title"]["runs"]),
+                    }
+                )
         return results
 
     @staticmethod
-    def _search(
-        query: Optional[str] = None, filter: SearchFilter = SearchFilter.VIDEO
-    ) -> Dict[str, Any]:
+    def _search(query: Optional[str] = None, filter: SearchFilter = SearchFilter.VIDEO):
         client = innertube.InnerTube("WEB", "2.20230920.00.00")
         response = client.search(query=query, params=filter.code if filter else None)
         return response
@@ -71,7 +64,7 @@ class SearchService:
         return f"https://www.youtube.com/watch?v={video_id}"
 
     @staticmethod
-    def get_stream_url(youtube_url):
+    def get_stream(youtube_url):
         try:
             session = streamlink.Streamlink()
             streams = session.streams(youtube_url)
@@ -79,27 +72,27 @@ class SearchService:
                 best_stream = streams.get("best")
                 return best_stream.url if best_stream else None
             else:
-                logging.warning("No streams found for this URL")
+                gr.Warning(f"No streams found for: {youtube_url}")
                 return None
         except Exception as e:
+            gr.Error(f"An error occurred while getting stream: {e}")
             logging.warning(f"An error occurred: {e}")
             return None
 
+
 INITIAL_STREAMS = SearchService.search("world live cams", SearchFilter.LIVE)
-class YouTubeObjectDetection:
+
+
+class LiveYouTubeObjectDetector:
     def __init__(self):
         logging.getLogger().setLevel(logging.DEBUG)
         self.model = YOLO("yolov8x.pt")
-        self.font_path = self.download_font(
-            "https://www.fontsquirrel.com/fonts/download/open-sans",
-            "open-sans.zip",
-        )
         self.current_page_token = None
         self.streams = INITIAL_STREAMS
 
-        # Gradio UI Elements
+        # Gradio UI
         initial_gallery_items = [(stream["thumbnail_url"], stream["title"]) for stream in self.streams]
-        self.gallery = gr.Gallery(label="Live YouTube Videos", value=initial_gallery_items, show_label=True, columns=[3], rows=[10], object_fit="contain", height="auto", allow_preview=False)
+        self.gallery = gr.Gallery(label="Live YouTube Videos", value=initial_gallery_items, show_label=True, columns=[4], rows=[5], object_fit="contain", height="auto", allow_preview=False)
         self.search_input = gr.Textbox(label="Search Live YouTube Videos")
         self.stream_input = gr.Textbox(label="URL of Live YouTube Video")
         self.annotated_image = gr.AnnotatedImage(show_label=False)
@@ -107,22 +100,16 @@ class YouTubeObjectDetection:
         self.submit_button = gr.Button("Detect Objects", variant="primary", size="lg")
         self.page_title = gr.HTML("<center><h1><b>Object Detection in Live YouTube Streams</b></h1></center>")
 
-
-    @staticmethod
-    def download_font(url, save_path):
-        os.system(f"wget {url} -O {save_path}")
-        with zipfile.ZipFile(save_path, "r") as zip_ref:
-            zip_ref.extractall(".")
-        return os.path.join(".", "OpenSans-Regular.ttf")
-
-    def capture_frame(self, url):
-        stream_url = SearchService.get_stream_url(url)
+    def detect_objects(self, url):
+        stream_url = SearchService.get_stream(url)
         if not stream_url:
-            return [], []
+            gr.Error(f"Unable to find a stream for: {stream_url}")
+            return self.create_black_image(), []
         frame = self.get_frame(stream_url)
         if frame is None:
-            return self.create_error_image("Failed to capture frame"), []
-        return self.process_frame(frame)
+            gr.Error(f"Unable to capture frame for: {stream_url}")
+            return self.create_black_image(), []
+        return self.annotate(frame)
 
     def get_frame(self, stream_url):
         if not stream_url:
@@ -140,7 +127,7 @@ class YouTubeObjectDetection:
             logging.warning(f"An error occurred while capturing the frame: {e}")
             return None
 
-    def process_frame(self, frame):
+    def annotate(self, frame):
         frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
         results = self.model(frame_rgb)
         annotations = self.get_annotations(results)
@@ -152,35 +139,21 @@ class YouTubeObjectDetection:
             for box in result.boxes:
                 class_id = int(box.cls[0])
                 class_name = result.names[class_id]
-                bbox = tuple(map(int, box.xyxy[0]))
-                annotations.append((bbox, class_name))
+                # EXTRACT BOUNDING BOX AND CONVERT TO INTEGER
+                x1, y1, x2, y2 = box.xyxy[0]
+                bbox_coords = (int(x1), int(y1), int(x2), int(y2))
+                annotations.append((bbox_coords, class_name))
         return annotations
 
-    def create_error_image(self, message):
-        error_image = np.zeros((1920, 1080, 3), dtype=np.uint8)
-        pil_image = Image.fromarray(error_image)
-        draw = ImageDraw.Draw(pil_image)
-        font = ImageFont.truetype(self.font_path, 24)
-        text_size = draw.textsize(message, font=font)
-        position = ((1920 - text_size[0]) // 2, (1080 - text_size[1]) // 2)
-        draw.text(position, message, (0, 0, 255), font=font)
-        return cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)
-
-    def fetch_live_streams(self, query=""):
-        streams = []
-        results = SearchService.search(query if query else "world live cams", SearchFilter.LIVE)
-        for result in results:
-            if "video_id" in result and "thumbnail_urls" in result:
-                streams.append(
-                    {
-                        "thumbnail_url": result["thumbnail_urls"][0]
-                        if result["thumbnail_urls"] else "",
-                        "title": result["title"],
-                        "video_id": result["video_id"],
-                        "label": result["video_id"],
-                    }
-                )
-        return streams
+    def create_black_image():
+        black_image = np.zeros((1080, 1920, 3), dtype=np.uint8)
+        pil_black_image = Image.fromarray(black_image)
+        cv2_black_image = cv2.cvtColor(np.array(pil_black_image), cv2.COLOR_RGB2BGR)
+        return cv2_black_image
+
+    @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}") as app:
@@ -199,26 +172,26 @@ class YouTubeObjectDetection:
                 self.gallery.render()
 
             @self.gallery.select(inputs=None, outputs=[self.annotated_image, self.stream_input])
-            def on_gallery_select(evt: gr.SelectData):
-                selected_index = evt.index
-                if selected_index is not None and selected_index < len(self.streams):
-                    selected_stream = self.streams[selected_index]
+            def detect_objects_from_gallery_item(evt: gr.SelectData):
+                if evt.index is not None and evt.index < len(self.streams):
+                    selected_stream = self.streams[evt.index]
                     stream_url = SearchService.get_youtube_url(selected_stream["video_id"])
-                    frame_output = self.capture_frame(stream_url)
+                    frame_output = self.detect_objects(stream_url)
                     return frame_output, stream_url
                 return None, ""
 
             @self.search_button.click(inputs=[self.search_input], outputs=[self.gallery])
-            def on_search_click(query):
-                self.streams = self.fetch_live_streams(query)
+            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 annotate_stream(url):
-                return self.capture_frame(url)
+            def detect_objects_from_url(url):
+                return self.detect_objects(url)
+
+        return app.queue().launch(show_api=False, debug=True, quiet=False, share=True)
 
-        return app.queue().launch(show_api=False, debug=True, quiet=False, share=False)
 
 if __name__ == "__main__":
-    YouTubeObjectDetection().render()
+    LiveYouTubeObjectDetector().render()