from ultralytics import YOLO
from PIL import Image
import gradio as gr
import yaml
import os
import numpy as np

# Define paths using relative paths for Hugging Face Space deployment
model_weights_path = 'best.pt'
data_yaml_path = 'data.yaml'

# Load the trained YOLOv8 model
model = YOLO(model_weights_path)

# Load the data.yaml file to extract class names
with open(data_yaml_path, 'r') as f:
    data_yaml_content = yaml.safe_load(f)
    class_names = data_yaml_content['names']

def detect_municipal_problems(image: Image.Image) -> Image.Image:
    """
    Performs object detection on an input image using the trained YOLOv8 model
    and returns the image with detected bounding boxes and labels.

    Args:
        image (PIL.Image.Image): The input image to analyze.

    Returns:
        PIL.Image.Image: The image with detected objects and bounding boxes.
    """
    print("Received image for detection.")

    # Perform prediction using the loaded model
    results = model.predict(source=image, imgsz=640, conf=0.25)

    # Assuming only one image is processed at a time
    if results:
        annotated_image_np = results[0].plot() # plot() returns an RGB numpy array
        # Convert the annotated NumPy array (RGB) back to PIL Image
        annotated_image = Image.fromarray(annotated_image_np)
        print("Detection complete. Image annotated.")
        return annotated_image
    else:
        print("No detections found or an issue occurred during prediction.")
        return image # Return original image if no detections or error

# Create the Gradio interface
interface = gr.Interface(
    fn=detect_municipal_problems,
    inputs=gr.Image(type='pil', label='Upload Image'),
    outputs=gr.Image(type='pil', label='Detected Problems'),
    title='Municipal Problem Detector using YOLOv8',
    description='Upload an image to detect municipal problems like Potholes, Flooding, and Waste Management.',
    live=True,
    examples=['/content/dataset2/test/images/image_002.jpg'] # Example image from test set
)

# Launch the Gradio application
if __name__ == '__main__':
    interface.launch(debug=True, share=True)