Update app.py

app.py

@@ -1,2 +1,283 @@
-import streamlit as st
-st.write("Hello, World!")
+import logging
+import os
+import sys
+import zipfile
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+os.system("python3 -m pip uninstall -y typing-extensions")
+os.system("python3 -m pip install -U typing-extensions")
+os.system(
+    "python3 -m pip install -q --progress-bar off streamlink gradio tiktoken ultralytics pillow innertube"
+)
+import cv2
+import gradio as gr
+import innertube
+import numpy as np
+import streamlink
+from PIL import Image, ImageDraw, ImageFont
+from ultralytics import YOLO
+
+logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)
+
+model = YOLO("yolov8x.pt")
+
+
+class SearchFilter(Enum):
+    LIVE = ("EgJAAQ%3D%3D", "Live")
+    VIDEO = ("EgIQAQ%3D%3D", "Video")
+
+    def __init__(self, code, human_readable):
+        self.code = code
+        self.human_readable = human_readable
+
+    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]):
+        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]]:
+        results = []
+        items = []
+
+        contents = (
+            data.get("contents", {})
+            .get("twoColumnSearchResultsRenderer", {})
+            .get("primaryContents", {})
+            .get("sectionListRenderer", {})
+            .get("contents", [])
+        )
+        if contents:
+            items = contents[0].get("itemSectionRenderer", {}).get("contents", [])
+
+        for item in items:
+            if "videoRenderer" in item:
+                renderer = item["videoRenderer"]
+                video_id = renderer.get("videoId", "")
+                thumbnail_urls = [
+                    thumb.get("url", "")
+                    for thumb in renderer.get("thumbnail", {}).get("thumbnails", [])
+                ]
+                title_text = "".join(
+                    [
+                        run.get("text", "")
+                        for run in renderer.get("title", {}).get("runs", [])
+                    ]
+                )
+
+                result = {
+                    "video_id": video_id,
+                    "thumbnail_urls": thumbnail_urls,
+                    "title": title_text,
+                }
+                results.append(result)
+
+        return results
+
+    @staticmethod
+    def _search(
+        query: Optional[str] = None, filter: SearchFilter = SearchFilter.VIDEO
+    ) -> Dict[str, Any]:
+        client = innertube.InnerTube("WEB", "2.20230920.00.00")
+        response = client.search(query=query, params=filter.code if filter else None)
+        return response
+
+    @staticmethod
+    def get_youtube_url(video_id: str) -> str:
+        return f"https://www.youtube.com/watch?v={video_id}"
+
+    @staticmethod
+    def get_stream_url(youtube_url):
+        try:
+            session = streamlink.Streamlink()
+            streams = session.streams(youtube_url)
+            if streams:
+                best_stream = streams.get("best")
+                return best_stream.url if best_stream else None
+            else:
+                logging.warning("No streams found for this URL")
+                return None
+        except Exception as e:
+            logging.warning(f"An error occurred: {e}")
+            return None
+
+
+class LiveStreamAnnotator:
+    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",
+            "/content/open-sans.zip",
+        )
+        self.current_page_token = None
+        self.streams = self.fetch_live_streams("world live cams")
+        # Gradio UI Elements
+        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=False,
+            columns=[3],
+            rows=[10],
+            object_fit="contain",
+            height="auto",
+        )
+        self.search_input = gr.Textbox(label="Search Live YouTube Videos")
+        self.stream_input = gr.Textbox(label="URL of Live YouTube Video")
+        self.output_image = gr.AnnotatedImage(show_label=False)
+        self.search_button = gr.Button("Search")
+        self.submit_button = gr.Button("Detect Objects", variant="primary", size="lg")
+        self.prev_page_button = gr.Button("Previous Page", interactive=False)
+        self.next_page_button = gr.Button("Next Page", interactive=False)
+
+    @staticmethod
+    def download_font(url, save_path):
+        os.system(f"wget {url} -O {save_path}")
+        with zipfile.ZipFile(save_path, "r") as zip_ref:
+            font_dir = "/usr/share/fonts/open-sans"
+            zip_ref.extractall(font_dir)
+        return os.path.join(font_dir, "OpenSans-Regular.ttf")
+
+    def capture_frame(self, url):
+        stream_url = SearchService.get_stream_url(url)
+        if not stream_url:
+            return self.create_error_image("No stream found"), []
+        frame = self.get_frame(stream_url)
+        if frame is None:
+            return self.create_error_image("Failed to capture frame"), []
+        return self.process_frame(frame)
+
+    def get_frame(self, stream_url):
+        if not stream_url:
+            return None
+        try:
+            cap = cv2.VideoCapture(stream_url)
+            ret, frame = cap.read()
+            cap.release()
+            if ret:
+                return cv2.resize(frame, (1920, 1080))
+            else:
+                logging.warning(
+                    "Unable to process the HLS stream with cv2.VideoCapture."
+                )
+                return None
+        except Exception as e:
+            logging.warning(f"An error occurred while capturing the frame: {e}")
+            return None
+
+    def process_frame(self, frame):
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        results = self.model(frame_rgb)
+        annotations = self.get_annotations(results)
+        return Image.fromarray(frame_rgb), annotations
+
+    def get_annotations(self, results):
+        annotations = []
+        for result in results:
+            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))
+        return annotations
+
+    @staticmethod
+    def create_error_image(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("/usr/share/fonts/open-sans/OpenSans-Regular.ttf", 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 render(self):
+        with gr.Blocks(
+            title="Object Detection in Live YouTube Streams",
+            css="footer {visibility: hidden}",
+        ) as app:
+            gr.HTML(
+                "<center><h1><b>Object Detection in Live YouTube Streams</b></h1></center>"
+            )
+            with gr.Column():
+                self.stream_input.render()
+                with gr.Group():
+                    self.output_image.render()
+                    self.submit_button.render()
+            with gr.Group():
+                with gr.Row():
+                    self.search_input.render()
+                    self.search_button.render()
+                with gr.Row():
+                    self.gallery.render()
+
+            @self.gallery.select(
+                inputs=None, outputs=[self.output_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]
+                    stream_url = SearchService.get_youtube_url(
+                        selected_stream["video_id"]
+                    )
+                    frame_output = self.capture_frame(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)
+                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.output_image]
+            )
+            def annotate_stream(url):
+                return self.capture_frame(url)
+
+        app.queue().launch(
+            share=True, debug=True, quiet=True, show_api=False, height=800
+        )
+
+
+LiveStreamAnnotator().render()