Blame for app.py

import os
import cv2
import torch
import tempfile
import requests
import base64
import mimetypes
import urllib.parse
import instaloader
from ultralytics import YOLO
import gradio as gr
import numpy as np
from PIL import Image
from pytube import YouTube
from io import BytesIO
from typing import Optional, Tuple, Union

def load_yolov8_model(model_name: str = 'yolov8s.pt') -> Optional[YOLO]:
    """
    Load a YOLOv8 model.

    Args:
        model_name (str): Name of the model to load.

    Returns:
        Optional[YOLO]: Loaded YOLO model or None if loading fails.
    """
    try:
        return YOLO(model_name)
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

def detect_objects(image: np.ndarray, model: YOLO) -> Optional[np.ndarray]:
    """
    Perform object detection on an image.

    Args:
        image (np.ndarray): Input image.
        model (YOLO): Loaded YOLO model.

    Returns:
        Optional[np.ndarray]: Processed image with bounding boxes or None if detection fails.
    """
    try:
        results = model(image)
        
        for result in results[0].boxes:
            x1, y1, x2, y2 = map(int, result.xyxy[0])
            label = model.names[int(result.cls)]
            confidence = result.conf.item()

            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label_text = f'{label} {confidence:.2f}'
            cv2.putText(image, label_text, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

        return image
    except Exception as e:
        print(f"Error during object detection: {e}")
        return None

def process_image_with_yolov8(model_name: str, image: Optional[Image.Image] = None, url: Optional[str] = None) -> Optional[np.ndarray]:
    """
    Process an image with YOLOv8 model.

    Args:
        model_name (str): Name of the YOLO model to use.
        image (Optional[Image.Image]): Input image.
        url (Optional[str]): URL of the image to process.

    Returns:
        Optional[np.ndarray]: Processed image or None if processing fails.
    """
    model = load_yolov8_model(model_name)
    if model is None:
        return None
    
    try:
        if url:
            response = requests.get(url)
            response.raise_for_status()
            image = Image.open(BytesIO(response.content))
        
        if image is None:
            raise ValueError("No image provided")
        
        return detect_objects(np.array(image), model)
    except Exception as e:
        print(f"Error processing image: {e}")
        return None

def download_image(url: str, save_path: str) -> None:
    """
    Download an image from a URL and save it to a specified path.

    Args:
        url (str): URL of the image to download.
        save_path (str): Path to save the downloaded image.
    """
    try:
        response = requests.get(url)
        response.raise_for_status()
        image = Image.open(BytesIO(response.content))
        image.save(save_path)
        print(f"Image saved to {save_path}")
    except requests.exceptions.RequestException as e:
        print(f"Failed to retrieve image. Error: {e}")

def download_base64_image(base64_data: str, save_path: str) -> None:
    """
    Decode a base64 image and save it to a specified path.

    Args:
        base64_data (str): Base64 encoded image data.
        save_path (str): Path to save the decoded image.
    """
    try:
        header, encoded = base64_data.split(',', 1)
        data = base64.b64decode(encoded)
        image = Image.open(BytesIO(data))
        image.save(save_path)
        print(f"Image saved to {save_path}")
    except Exception as e:
        print(f"Failed to decode base64 image. Error: {e}")

def download_video(url: str, save_path: str) -> None:
    """
    Download a video from a URL and save it to a specified path.

    Args:
        url (str): URL of the video to download.
        save_path (str): Path to save the downloaded video.
    """
    try:
        yt = YouTube(url)
        stream = yt.streams.get_highest_resolution()
        stream.download(output_path=save_path)
        print(f"Video saved to {os.path.join(save_path, stream.default_filename)}")
    except Exception as e:
        print(f"Failed to retrieve video. Error: {e}")

def download_instagram_post(url: str, save_path: str) -> None:
    """
    Download an Instagram post (image or video) and save it to a specified path.

    Args:
        url (str): URL of the Instagram post.
        save_path (str): Path to save the downloaded content.
    """
    loader = instaloader.Instaloader()
    post = instaloader.Post.from_shortcode(loader.context, url.split("/")[-2])
    if post.is_video:
        loader.download_post(post, target=save_path)
        print(f"Video saved to {save_path}")
    else:
        for file in post.get_sidecar_nodes():
            file_url = file.display_url
            download_image(file_url, os.path.join(save_path, f'insta_image_{file.shortcode}.jpg'))
        print(f"Images saved to {save_path}")

def download_youtube_video(youtube_url: str) -> Optional[str]:
    """
    Download a YouTube video and return the path to the downloaded file.

    Args:
        youtube_url (str): URL of the YouTube video.

    Returns:
        Optional[str]: Path to the downloaded video file or None if download fails.
    """
    try:
        yt = YouTube(youtube_url)
        stream = yt.streams.filter(file_extension='mp4', res="720p").first()
        if not stream:
            stream = yt.streams.filter(file_extension='mp4').first()
        output_path = os.path.join(tempfile.gettempdir(), "downloaded_video.mp4")
        stream.download(output_path=os.path.dirname(output_path), filename=os.path.basename(output_path))
        return output_path
    except Exception as e:
        print(f"Failed to retrieve video from YouTube. Error: {e}")
        return None

def process_video(input_video_path: str, output_video_path: str, model: YOLO) -> Optional[str]:
    """
    Process a video file with object detection and save the result.

    Args:
        input_video_path (str): Path to the input video file.
        output_video_path (str): Path to save the processed video.
        model (YOLO): Loaded YOLO model.

    Returns:
        Optional[str]: Path to the processed video or None if processing fails.
    """
    cap = cv2.VideoCapture(input_video_path)
    if not cap.isOpened():
        print(f"Error: Cannot open video file {input_video_path}")
        return None

    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')

    out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))

    while True:
        ret, frame = cap.read()
        if not ret:
            break

        processed_frame = detect_objects(frame, model)
        out.write(processed_frame)

    cap.release()
    out.release()

    return output_video_path

def handle_input_selection(model_name: str, choice: str, image: Optional[Image.Image], url: Optional[str]) -> Tuple[Optional[np.ndarray], Optional[str]]:
    """
    Handle input selection and process the image accordingly.

    Args:
        model_name (str): Name of the YOLO model to use.
        choice (str): Input method choice ('Upload' or 'URL').
        image (Optional[Image.Image]): Uploaded image.
        url (Optional[str]): Image URL.

    Returns:
        Tuple[Optional[np.ndarray], Optional[str]]: Processed image and error message (if any).
    """
    if choice == "Upload" and image is not None:
        return process_image_with_yolov8(model_name, image=image), None
    elif choice == "URL" and url:
        return process_image_with_yolov8(model_name, url=url), None
    else:
        return None, "Invalid input. Please provide an image or URL."

def process_input(model_name: str, local_video: Optional[str] = None, video_url: Optional[str] = None) -> Union[str, Optional[str]]:
    """
    Process input video with YOLOv8 model.

    Args:
        model_name (str): Name of the YOLO model to use.
        local_video (Optional[str]): Path to local video file.
        video_url (Optional[str]): YouTube video URL.

    Returns:
        Union[str, Optional[str]]: Path to processed video or error message.
    """
    model = load_yolov8_model(model_name)
    if model is None:
        return "Failed to load YOLO model."
    
    output_path = os.path.join(tempfile.gettempdir(), "processed_video.mp4")
    
    if local_video:
        input_path = local_video
    elif video_url:
        input_path = download_youtube_video(video_url)
    else:
        return "Please provide a video file or YouTube URL."
    
    if input_path is None:
        return "Error downloading or accessing video."
    
    result = process_video(input_path, output_path, model)
    return result if result else "Error processing video."

# Create Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("## Multimedia Object Detection Using YOLO")
    
    model_choice = gr.Radio(["yolov8s.pt", "yolov8m.pt"], label="Select Model")
    
    with gr.Tabs() as tabs:
        with gr.TabItem("Image Detection"):
            input_choice = gr.Radio(["Upload", "URL"], label="Select Input Method")
            with gr.Row():
                with gr.Column(visible=False) as upload_section:
                    input_image = gr.Image(type="pil", label="Upload Image")
                with gr.Column(visible=False) as url_section:
                    input_url = gr.Textbox(label="Image URL", placeholder="Enter URL of the image")
            
            def toggle_inputs(choice):
                return (
                    gr.update(visible=choice == "Upload"),
                    gr.update(visible=choice == "URL")
                )
            
            input_choice.change(toggle_inputs, input_choice, [upload_section, url_section])
            output_image = gr.Image(type="numpy", label="Processed Image")
            process_button_image = gr.Button("Process Image")
            process_button_image.click(handle_input_selection, [model_choice, input_choice, input_image, input_url], [output_image, gr.State()])
        
        with gr.TabItem("Video Detection"):
            local_video = gr.Video(label="Upload Local Video")
            video_url = gr.Textbox(label="YouTube Video URL")
            output_video = gr.Video(label="Processed Video")
            process_button_video = gr.Button("Process Video")
            process_button_video.click(process_input, [model_choice, local_video, video_url], output_video)

    gr.Markdown("Upload an image, provide an image URL, or provide a YouTube video URL to perform object detection using YOLOv8 model.")
    
demo.launch()