app.py

"""
This module integrates real-time object detection into live YouTube streams using the YOLO (You Only Look Once) 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 Streamlink 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.

The module comprises several key components:
- `SearchService`: A service class to search for YouTube videos and retrieve live stream URLs.
- `LiveYouTubeObjectDetector`: The main class integrating the YOLO model and Gradio UI, handling the entire workflow of searching, streaming, and object detection.

Dependencies:
- OpenCV (`cv2`): Used for image processing tasks.
- Gradio: Provides the interactive web-based user interface.
- Streamlink: Used for retrieving live stream data.
- NumPy: Utilized for numerical operations on image data.
- Pillow (`PIL`): A Python Imaging Library for opening, manipulating, and saving images.
- Ultralytics YOLO: The YOLO model implementation for object detection.
- `youtube-search-python`: Used for searching YouTube without API keys.
- `imageio`: For reading frames from live streams using FFmpeg.

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

import asyncio
import cv2
import gradio as gr
import numpy as np
from ultralytics import YOLO
import streamlink
from PIL import Image
from youtubesearchpython import CustomSearch, SearchMode
import imageio.v3 as iio

logging.basicConfig(level=logging.DEBUG)


class SearchService:
    """
    SearchService provides functionality to search for YouTube videos using the
    `youtube-search-python` library and retrieve live stream URLs using the Streamlink library.

    Methods:
        search: Searches YouTube for videos matching a query and live filter.
        get_youtube_url: Constructs a YouTube URL for a given video ID.
        get_stream: Retrieves the stream URL for a given YouTube video URL.
    """

    @staticmethod
    def search(query: str, live: bool = False) -> List[Dict[str, Any]]:
        """
        Searches YouTube for videos matching the given query and live filter.

        :param query: The search query.
        :type query: str
        :param live: Whether to filter for live videos.
        :type live: bool
        :return: A list of search results, each a dictionary with video details.
        :rtype: List[Dict[str, Any]]
        """
        results = []
        # Apply live filter if needed
        if live:
            # Use CustomSearch with SearchMode.live to filter live videos
            search = CustomSearch(query, SearchMode.live, limit=20)
        else:
            # Use CustomSearch with default mode for regular videos
            search = CustomSearch(query, SearchMode.videos, limit=20)
        for result in search.result()['result']:
            if 'video' in result:
                video = result['video']
                results.append({
                    'video_id': video['id'],
                    'thumbnail_url': video['thumbnails'][-1]['url'],
                    'title': video['title'],
                })
        return results

    @staticmethod
    def get_youtube_url(video_id: str) -> str:
        """
        Constructs a YouTube URL for the given video ID.

        :param video_id: The ID of the YouTube video.
        :type video_id: str
        :return: The YouTube URL for the video.
        :rtype: 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.

        :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:
            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(f"No streams found for: {youtube_url}")
                return None
        except Exception as e:
            logging.warning(f"An error occurred while getting stream: {e}")
            return None


class LiveYouTubeObjectDetector:
    """
    LiveYouTubeObjectDetector is a class that integrates object detection into live YouTube streams.
    It uses the YOLO (You Only Look Once) model to detect objects in video frames captured from live streams.
    The class also provides a Gradio interface for users to interact with the object detection system,
    allowing them to search for live streams, view them, and detect objects in real-time.

    Methods:
        detect_objects: Detects objects in a live YouTube stream given its URL.
        get_frame: Captures a frame from a live stream URL.
        annotate: Annotates a frame with detected objects.
        create_black_image: Creates a black placeholder image.
        get_live_streams: Searches for live streams based on a query.
        render: Sets up and launches the Gradio interface.
    """

    def __init__(self):
        """Initializes the LiveYouTubeObjectDetector with YOLO model and UI components."""
        logging.getLogger().setLevel(logging.DEBUG)
        self.model = YOLO("yolov8x.pt")
        self.model.fuse()
        self.streams = self.get_live_streams("world live cams")

    async def detect_objects(self, url: str) -> Tuple[Image.Image, List[Tuple[Tuple[int, int, int, int], str]]]:
        """
        Detects objects in the given live YouTube stream URL.

        :param url: The URL of the live YouTube video.
        :type url: str
        :return: A tuple containing the annotated image and a list of annotations.
        :rtype: 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 = await 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)

    async def get_frame(self, stream_url: str) -> Optional[np.ndarray]:
        """
        Captures a frame from the given live stream URL.

        :param stream_url: The URL of the live stream.
        :type stream_url: str
        :return: The captured frame as a numpy array, or None if capture fails.
        :rtype: Optional[np.ndarray]
        """
        if not stream_url:
            return None
        try:
            reader = iio.imiter(stream_url, plugin='ffmpeg', fps=1)
            loop = asyncio.get_event_loop()
            frame = await loop.run_in_executor(None, next, reader, None)
            if frame is None:
                return None
            # Resize frame if needed
            frame = cv2.resize(frame, (1280, 720))
            return frame
        except StopIteration:
            logging.warning("Could not read frame from 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]]]:
        """
        Annotates the given frame with detected objects and their bounding boxes.

        :param frame: The frame to be annotated.
        :type frame: np.ndarray
        :return: A tuple of the annotated PIL image and list of annotations.
        :rtype: Tuple[Image.Image, List[Tuple[Tuple[int, int, int, int], str]]]
        """
        results = self.model(frame)[0]
        annotations = []
        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))
        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        return pil_image, annotations

    @staticmethod
    def create_black_image() -> Tuple[Image.Image, List]:
        """
        Creates a black image of fixed dimensions.

        :return: A black image as a PIL Image and an empty list of annotations.
        :rtype: Tuple[Image.Image, List]
        """
        black_image = np.zeros((720, 1280, 3), dtype=np.uint8)
        pil_black_image = Image.fromarray(black_image)
        return pil_black_image, []

    def get_live_streams(self, query: str = "") -> List[Dict[str, Any]]:
        """
        Searches for live streams on YouTube based on the given query.

        :param query: The search query for live streams.
        :type query: str
        :return: A list of dictionaries containing information about each live stream.
        :rtype: List[Dict[str, str]]
        """
        return SearchService.search(query if query else "world live cams", live=True)

    def render(self):
        """
        Sets up and launches the Gradio interface for the application.

        This method creates the Gradio UI elements and defines the behavior of the application.
        It includes the setup of interactive widgets like galleries, textboxes, and buttons,
        and defines the actions triggered by user interactions with these widgets.

        The Gradio interface allows users to search for live YouTube streams, select a stream,
        and run object detection on the selected live stream.
        """
        with gr.Blocks(title="Object Detection in Live YouTube Streams", css=".gradio-container {background-color: #f5f5f5}", theme=gr.themes.Soft()) as app:
            gr.HTML("<h1 style='text-align: center; color: #1E88E5;'>Object Detection in Live YouTube Streams</h1>")
            with gr.Tabs():
                with gr.TabItem("Live Stream Detection"):
                    with gr.Row():
                        stream_input = gr.Textbox(label="URL of Live YouTube Video", placeholder="Enter YouTube live stream URL...", interactive=True)
                        submit_button = gr.Button("Detect Objects", variant="primary")
                    annotated_image = gr.AnnotatedImage(label="Detection Result", height=480)
                    status_text = gr.Markdown(value="Status: Ready", visible=False)

                    async def detect_objects_from_url(url):
                        status_text.update(value="Status: Processing...", visible=True)
                        try:
                            result = await self.detect_objects(url)
                            status_text.update(value="Status: Done", visible=True)
                            return result
                        except Exception as e:
                            logging.error(f"An error occurred: {e}")
                            status_text.update(value=f"Status: Error - {e}", visible=True)
                            return self.create_black_image()

                    submit_button.click(fn=detect_objects_from_url, inputs=[stream_input], outputs=[annotated_image], api_name="detect_objects")

                with gr.TabItem("Search Live Streams"):
                    with gr.Row():
                        search_input = gr.Textbox(label="Search for Live YouTube Streams", placeholder="Enter search query...", interactive=True)
                        search_button = gr.Button("Search", variant="secondary")
                    gallery = gr.Gallery(label="Live YouTube Streams", show_label=False).style(grid=[4], height="auto")
                    gallery.style(item_height=150)
                    status_text_search = gr.Markdown(value="", visible=False)

                    def search_live_streams(query):
                        status_text_search.update(value="Searching...", visible=True)
                        self.streams = self.get_live_streams(query)
                        gallery_items = []
                        for stream in self.streams:
                            thumb_url = stream["thumbnail_url"]
                            title = stream["title"]
                            video_id = stream["video_id"]
                            gallery_items.append((thumb_url, title, video_id))
                        status_text_search.update(value="Search Results:", visible=True)
                        return gr.update(value=gallery_items)

                    search_button.click(fn=search_live_streams, inputs=[search_input], outputs=[gallery], api_name="search_streams")

                    async def detect_objects_from_gallery_item(evt: gr.SelectData):
                        index = evt.index
                        if index is not None and index < len(self.streams):
                            selected_stream = self.streams[index]
                            stream_url = SearchService.get_youtube_url(selected_stream["video_id"])
                            stream_input.value = stream_url
                            result = await self.detect_objects(stream_url)
                            annotated_image.update(value=result[0], annotations=result[1])
                            with gr.Row():
                                annotated_image.render()
                            return result

                    gallery.select(fn=detect_objects_from_gallery_item, inputs=None, outputs=None)

            gr.Markdown(
                """
                **Instructions:**
                - **Live Stream Detection Tab:** Enter a YouTube live stream URL and click 'Detect Objects' to view the real-time object detection.
                - **Search Live Streams Tab:** Search for live streams on YouTube, select one from the gallery, and view object detection results.
                """
            )

            gr.HTML("<footer style='text-align: center; color: gray;'>Developed using Gradio and YOLO</footer>")

        app.queue(concurrency_count=3).launch()

if __name__ == "__main__":
    LiveYouTubeObjectDetector().render()