# Import necessary libraries
import gradio as gr  # Gradio for creating the web interface
import cv2  # OpenCV for image processing
from huggingface_hub import hf_hub_download  # Download models from Hugging Face Hub
from gradio_webrtc import WebRTC  # WebRTC integration for streaming webcam feeds in Gradio
from twilio.rest import Client  # Twilio client for managing ICE servers for WebRTC
import os  # OS module for environment variable access
from inference import YOLOv10  # Custom YOLOv10 inference class

# Download YOLOv10 model file from Hugging Face Hub
model_file = hf_hub_download(
    repo_id="onnx-community/yolov10n",  # Repository containing the YOLOv10 model
    filename="onnx/model.onnx"  # Model file to download
)

# Initialize the YOLOv10 model
model = YOLOv10(model_file)

# Retrieve Twilio account credentials from environment variables
account_sid = os.environ.get("TWILIO_ACCOUNT_SID")  # Twilio Account SID
auth_token = os.environ.get("TWILIO_AUTH_TOKEN")  # Twilio Auth Token

# Check if Twilio credentials are available
if account_sid and auth_token:
    # Initialize Twilio client with credentials
    client = Client(account_sid, auth_token)

    # Create a Twilio token for ICE server configuration
    token = client.tokens.create()

    # Configure WebRTC to use Twilio ICE servers for better connection reliability
    rtc_configuration = {
        "iceServers": token.ice_servers,  # Use Twilio ICE servers
        "iceTransportPolicy": "relay",  # Relay policy to improve NAT traversal
    }
else:
    # Use default WebRTC configuration if Twilio credentials are not available
    rtc_configuration = None

# Function to perform object detection
def detection(image, conf_threshold=0.3):
    # Resize the input image to match the model's expected dimensions
    image = cv2.resize(image, (model.input_width, model.input_height))
    # Perform object detection and return the processed image
    new_image = model.detect_objects(image, conf_threshold)
    return cv2.resize(new_image, (500, 500))  # Resize output image for display

# Define custom CSS for Gradio interface layout
css = """.my-group {max-width: 600px !important; max-height: 600 !important;}
                      .my-column {display: flex !important; justify-content: center !important; align-items: center !important};"""

# Create a Gradio interface with custom blocks
with gr.Blocks(css=css) as demo:
    # Add a title to the Gradio interface
    gr.HTML(
        """
    <h1 style='text-align: center'>
    YOLOv10 Webcam Stream
    </h1>
    """
    )
    # Add links to the arXiv paper and GitHub repository for YOLOv10
    gr.HTML(
        """
        <h3 style='text-align: center'>
        <a href='https://arxiv.org/abs/2405.14458' target='_blank'>arXiv</a> | <a href='https://github.com/THU-MIG/yolov10' target='_blank'>github</a>
        </h3>
        """
    )
    # Define a column layout for the interface
    with gr.Column(elem_classes=["my-column"]):
        with gr.Group(elem_classes=["my-group"]):
            # Add a WebRTC component for webcam streaming
            image = WebRTC(label="Stream", rtc_configuration=rtc_configuration)
            # Add a slider to adjust the confidence threshold for object detection
            conf_threshold = gr.Slider(
                label="Confidence Threshold",  # Label for the slider
                minimum=0.0,  # Minimum slider value
                maximum=1.0,  # Maximum slider value
                step=0.05,  # Step size for slider
                value=0.30,  # Default slider value
            )

        # Stream webcam frames through the detection function
        image.stream(
            fn=detection,  # Detection function to process frames
            inputs=[image, conf_threshold],  # Inputs: webcam stream and confidence threshold
            outputs=[image],  # Outputs: processed frames
            time_limit=10  # Limit each detection to 10 seconds
        )

# Launch the Gradio app when the script is run directly
if __name__ == "__main__":
    demo.launch()