import os
os.environ['YOLO_CONFIG_DIR'] = '/home/user/.config/Ultralytics'
from ultralytics import YOLO
from PIL import Image
import numpy as np
import io
import gradio as gr
import cv2

# Load the model with the saved weights on CPU
model = YOLO('best.pt').to('cpu')

def process_webcam(image):
    # Perform inference
    results = model.predict(source=image)
    
    # Get the annotated image with bounding boxes
    annotated_image = results[0].plot()
    
    # Convert annotated image to PIL Image for Gradio
    annotated_image_pil = Image.fromarray(annotated_image)
    
    # Process results for Gradio display
    processed_results = []
    for result in results:
        boxes = result.boxes.xyxy.tolist()
        classes = result.boxes.cls.tolist()
        confs = result.boxes.conf.tolist()
        
        for box, cls, conf in zip(boxes, classes, confs):
            processed_results.append({
                "box": box,
                "class": int(cls),
                "confidence": float(conf)
            })
    
    return annotated_image_pil, processed_results

def process_video(video):
    cap = cv2.VideoCapture(video)
    frames = []
    processed_results = []

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Convert frame to RGB
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # Perform inference
        results = model.predict(source=frame_rgb)
        
        # Get the annotated image with bounding boxes
        annotated_frame = results[0].plot()
        
        # Append the annotated frame to the list
        frames.append(annotated_frame)
        
        # Process results for Gradio display
        for result in results:
            boxes = result.boxes.xyxy.tolist()
            classes = result.boxes.cls.tolist()
            confs = result.boxes.conf.tolist()
            
            for box, cls, conf in zip(boxes, classes, confs):
                processed_results.append({
                    "box": box,
                    "class": int(cls),
                    "confidence": float(conf)
                })

    cap.release()
    
    # Convert frames to video
    if frames:
        height, width, layers = frames[0].shape
        video_buffer = io.BytesIO()
        out = cv2.VideoWriter(video_buffer.name, cv2.VideoWriter_fourcc(*'mp4v'), 20.0, (width, height))

        for frame in frames:
            out.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
        
        out.release()
        video_buffer.seek(0)
        return video_buffer, processed_results
    else:
        return None, processed_results

def process_image(image):
    # Perform inference
    results = model.predict(source=image)
    
    # Get the annotated image with bounding boxes
    annotated_image = results[0].plot()
    
    # Convert annotated image to PIL Image for Gradio
    annotated_image_pil = Image.fromarray(annotated_image)
    
    # Process results for Gradio display
    processed_results = []
    for result in results:
        boxes = result.boxes.xyxy.tolist()
        classes = result.boxes.cls.tolist()
        confs = result.boxes.conf.tolist()
        
        for box, cls, conf in zip(boxes, classes, confs):
            processed_results.append({
                "box": box,
                "class": int(cls),
                "confidence": float(conf)
            })
    
    return annotated_image_pil, processed_results

# Define the Gradio interface
with gr.Blocks() as interface:
    gr.Markdown("# Object Detection with YOLO")
    
    with gr.Row():
        input_type = gr.Radio(["Live Webcam", "Video", "Image"], label="Input Type", value="Image")
    
    with gr.Row():
        with gr.Column():
            live_webcam = gr.Image(sources="webcam", streaming=True, label="Live Webcam", visible=False)
            video_input = gr.Video(label="Video Input", visible=False)
            image_input = gr.Image(type="numpy", label="Image Input")
        
        with gr.Column():
            live_output = gr.Image(label="Live Detection", visible=False)
            video_output = gr.Video(label="Processed Video", visible=False)
            image_output = gr.Image(label="Processed Image")
    
    json_output = gr.JSON(label="Detection Results")
    
    submit_button = gr.Button("Submit", visible=True)
    
    def update_input_type(choice):
        return {
            live_webcam: gr.update(visible=choice == "Live Webcam"),
            video_input: gr.update(visible=choice == "Video"),
            image_input: gr.update(visible=choice == "Image"),
            live_output: gr.update(visible=choice == "Live Webcam"),
            video_output: gr.update(visible=choice == "Video"),
            image_output: gr.update(visible=choice != "Live Webcam"),
            submit_button: gr.update(visible=choice != "Live Webcam")
        }
    
    input_type.change(update_input_type, input_type, 
                      [live_webcam, video_input, image_input, 
                       live_output, video_output, image_output, submit_button])
    
    live_webcam.stream(process_webcam, live_webcam, [live_output, json_output])
    
    def process_input(input_type, video, image):
        if input_type == "Video":
            video_buffer, results = process_video(video)
            return video_buffer, None, results
        else:
            image_pil, results = process_image(image)
            return None, image_pil, results
    
    submit_button.click(
        process_input,
        inputs=[input_type, video_input, image_input],
        outputs=[video_output, image_output, json_output]
    )

# Launch the Gradio 
interface.launch()