import gradio as gr
import requests
import numpy as np
from PIL import Image
import io
from ultralytics import YOLO
import cv2  

# Load the YOLO model at startup
try:
    model = YOLO('modelo_epoch_50.pt')
    print("Model loaded successfully")
except Exception as e:
    print(f"Error loading model: {str(e)}")
    model = None

# Function to process the captured image
def process_image(image):    
    # Check if image is None
    if image is None:
        raise gr.Error("Please take a picture first before analyzing!")
    
    # Convert image to RGB if it's not already
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Run inference
    results = model.predict(image, save=True, conf=0.5)

    # Print the results
    print("Model predictions:", results[0].boxes)

    # Convert PIL Image to numpy array for OpenCV
    image_cv = np.array(image)
    # Convert RGB to BGR (OpenCV uses BGR format)
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)

    # Draw bounding boxes of the prediction
    boxes = results[0].boxes
    for box in boxes:
        b = box.xyxy[0]  # Bounding box coordinates
        c = box.cls  # Predicted class
        confidence = box.conf

        x1, y1, x2, y2 = map(int, b)
        cv2.rectangle(image_cv, (x1, y1), (x2, y2), (255, 0, 0), 2)  # Blue for prediction

        label = f"{results[0].names[int(c)]} {confidence.item():.2f}"
        cv2.putText(image_cv, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
        print(label)

    # Convert back to RGB for display if needed
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
    # Convert back to PIL Image
    final_image = Image.fromarray(image_cv)
    # Save the final image to current directory, overwriting if exists
    final_image.save('FINAL.jpg', 'JPEG', quality=95)
    
    return final_image

# Create the Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        # Left column for camera and controls
        with gr.Column(scale=1):
            gr.Markdown("### Step 1: Take a Picture")
            camera = gr.Image(type="pil", label="Camera View", sources=["webcam"])
            gr.Markdown("Click the '📸 Take Photo' button in the camera view above")
            
            with gr.Row():
                analyze_btn = gr.Button("🔍 Analyze", variant="primary")
                new_picture_btn = gr.Button("🔄 Reset")
        
        # Right column for results
        with gr.Column(scale=1):
            gr.Markdown("### Step 2: View Results")
            modified_image = gr.Image(label="Analyzed Image")
    
    # Set up the event handlers
    analyze_btn.click(
        fn=process_image,
        inputs=[camera],
        outputs=[modified_image]
    )
    
    new_picture_btn.click(
        fn=lambda: (None, None),
        inputs=[],
        outputs=[camera, modified_image]
    )

# Launch the interface
if __name__ == "__main__":
    demo.launch(share=True)