app.py

"""
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.
- Retrieve live stream URLs using the `pytube` library.
- Perform real-time object detection on live streams using the YOLO model.
- Display the live stream and object detection results through a Gradio interface.

Dependencies:
- cv2 (OpenCV): Used for image processing tasks.
- Gradio: Provides the interactive web-based user interface.
- `pytube`: Used for retrieving live stream URLs from YouTube.
- innertube: Used for interacting with YouTube's internal API.
- numpy: Utilized for numerical operations on image data.
- PIL (Pillow): A Python Imaging Library for opening, manipulating, and saving images.
- ultralytics YOLO: The YOLO model implementation for object detection.

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
from typing import Any, Dict, List, Optional, Tuple

import cv2
import gradio as gr
import innertube
import numpy as np
from PIL import Image
from ultralytics import YOLO
from pytube import YouTube

logging.basicConfig(stream=sys.stderr, level=logging.DEBUG)


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):
        response = SearchService._search(query, filter)
        results = SearchService.parse(response)
        return results

    @staticmethod
    def parse(data: Dict[str, Any]) -> List[Dict[str, str]]:
        results = []
        try:
            contents = data["contents"]["twoColumnSearchResultsRenderer"]["primaryContents"]["sectionListRenderer"]["contents"]
            for content in contents:
                items = content.get("itemSectionRenderer", {}).get("contents", [])
                for item in items:
                    if "videoRenderer" in item:
                        renderer = item["videoRenderer"]
                        video_id = renderer.get("videoId", "")
                        thumbnails = renderer.get("thumbnail", {}).get("thumbnails", [])
                        thumbnail_url = thumbnails[-1]["url"] if thumbnails else ""
                        title_runs = renderer.get("title", {}).get("runs", [])
                        title = "".join(run.get("text", "") for run in title_runs)
                        results.append(
                            {
                                "video_id": video_id,
                                "thumbnail_url": thumbnail_url,
                                "title": title,
                            }
                        )
        except Exception as e:
            logging.error(f"Error parsing search results: {e}")
        return results

    @staticmethod
    def _search(query: Optional[str] = None, filter: SearchFilter = SearchFilter.VIDEO) -> Dict[str, Any]:
        client = innertube.InnerTube(client_name="WEB", client_version="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(youtube_url: str) -> Optional[str]:
        """Retrieves the 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.
        :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
            else:
                logging.warning(f"No live stream found for: {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)
        self.model = YOLO("yolo11n.pt")
        self.streams = INITIAL_STREAMS

        # 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=[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)
        self.search_button = gr.Button("Search", size="lg")
        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>")

    def detect_objects(self, url: str) -> Tuple[Image.Image, List[Tuple[Tuple[int, int, int, int], str]]]:
        stream_url = SearchService.get_stream(url)
        if not stream_url:
            logging.error(f"Unable to find a stream for: {url}")
            return self.create_black_image()
        frame = self.get_frame(stream_url)
        if frame is None:
            logging.error(f"Unable to capture frame for: {url}")
            return self.create_black_image()
        return self.annotate(frame)

    def get_frame(self, stream_url: str) -> Optional[np.ndarray]:
        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()
            if ret and frame is not None:
                return cv2.resize(frame, (1920, 1080))
            else:
                logging.warning("Unable to process the live stream with cv2.VideoCapture.")
                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)
        annotations = []
        result = predictions[0]

        boxes = result.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

    @staticmethod
    def create_black_image() -> Tuple[Image.Image, List]:
        black_image = np.zeros((1080, 1920, 3), dtype=np.uint8)
        pil_black_image = Image.fromarray(black_image)
        return pil_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}", analytics_enabled=False) as app:
            self.page_title.render()
            with gr.Column():
                with gr.Group():
                    with gr.Row():
                        self.stream_input.render()
                        self.submit_button.render()
                self.annotated_image.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.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):
                    selected_stream = self.streams[evt.index]
                    stream_url = SearchService.get_youtube_url(selected_stream["video_id"])
                    self.stream_input.value = stream_url
                    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()