import gradio as gr
from ultralytics import YOLO
from PIL import Image

# --- Documentation Strings ---

USAGE_GUIDELINES = """
## 1. Quick Start Guide: Processing an Image
This application uses a specialized YOLO model for object detection to automatically locate and classify specific "Port" features within an image.

1.  **Upload Image**: Click the 'Upload Image' box and select your image file (JPG or PNG).
2.  **Run**: Click the **"Process Image"** button.
3.  **Review**: The 'Output Image with labels' will display the detected ports, marked by bounding boxes, class labels, and confidence scores.
"""

INPUT_EXPLANATION = """
## 2. Expected Inputs and Best Practices

| Input Field | Purpose | Requirement |
| :--- | :--- | :--- |
| **Upload Image** | The image containing the ports or features you wish to classify. | Must be a standard image file (JPG, PNG). |

### Optimal Image Conditions
*   **Clarity and Focus:** Ensure the ports are clearly visible and in focus. Blurry images significantly reduce detection accuracy.
*   **Lighting:** Use well-lit images. Shadows or low-light conditions can cause the model to miss detections or misclassify objects.
*   **Framing:** The ports should occupy a significant portion of the image. Highly zoomed-out images may treat ports as background noise.
"""

OUTPUT_EXPLANATION = """
## 3. Expected Output 

The output image features graphical annotations provided by the AI model. Each annotation consists of three key components:

| Annotation Component | Description | Importance |
| :--- | :--- | :--- |
| **Bounding Box** | A colored rectangle drawn tightly around the detected object (the port). | Visually indicates the exact location of the object. |
| **Class Label** | The name of the predicted port type (e.g., USB-A, HDMI, Ethernet). | Represents the model's classification of the object within the box. |
| **Confidence Score** | A percentage (e.g., 0.92) representing the model's certainty in its prediction. | Scores below 0.50 (50%) should typically be treated as unreliable detections. |

"""

TECHNICAL_TIPS = """
## 4. Troubleshooting and Tips

*   **No Detections:** If the output image is identical to the input with no boxes, the model either failed to find any objects or the confidence for all potential detections was too low to be displayed. Try improving the lighting or clarity of the input photo.
*   **Testing with Examples:** Use the provided example images to quickly verify the application's functionality and see what types of objects the model is trained to recognize.
*   **Processing Time:** Processing time depends on the image resolution and the complexity of the scene. Large images require more computational resources for detection.
"""

# --------------------
# Core Pipeline Functions (Kept AS IS)
# --------------------

# Load the YOLO model weights
# Ensure 'best.pt' is available in the run directory
try:
    model = YOLO("./best.pt")
except Exception as e:
    print(f"Error loading model weights: {e}. Using a placeholder for demonstration.")
    # Define a placeholder function if model loading fails
    def placeholder_predict(img):
        return Image.new('RGB', img.size, color='red')
    model = placeholder_predict


def process_img(img: Image.Image):
    if img is None:
        gr.Warning("Please upload an image before processing.")
        return None
        

    # Perform prediction
    result = model.predict(img)
    
    # r.plot() returns a BGR numpy array
    for r in result:
        im_bgr = r.plot()
        # Convert BGR (OpenCV standard) to RGB (PIL/Gradio standard)
        return Image.fromarray(im_bgr[..., ::-1])


# --------------------
# Gradio UI
# --------------------

with gr.Blocks(title="Port Classification App") as demo:
    gr.Markdown("<h1 style='text-align: center;'> Port Classification and Object Detection </h1>")
    
    # 1. Guidelines Accordion
    with gr.Accordion("Tips & User Guidelines", open=False):
        gr.Markdown(USAGE_GUIDELINES)
        gr.Markdown("---")
        gr.Markdown(INPUT_EXPLANATION)
        gr.Markdown("---")
        gr.Markdown(OUTPUT_EXPLANATION)
        gr.Markdown("---")
        gr.Markdown(TECHNICAL_TIPS)

    # 2. Interface Definition
    with gr.Row():
        with gr.Column():
            gr.Markdown("## Step 1: Upload Port Image ")
            upload_img = gr.Image(label=" Upload Image", type="pil")
            gr.Markdown("## Step 2: Click Process Image ")
            classify_img_button = gr.Button(value=" Process Image", variant="primary")

        with gr.Column():
            gr.Markdown("## Result ")
            output_img = gr.Image(label=" Output Image with Detected Ports")
            
    # 3. Examples Section
    gr.Markdown("## Examples ")
    with gr.Row():
        gr.Examples(
            examples=[
                ["./examples/01.jpg"],
                ["./examples/02.jpg"],
                ["./examples/03.jpg"],
                ["./examples/04.jpg"],
                ["./examples/05.jpg"],
                ["./examples/06.jpg"],
            ],
            inputs=upload_img,
            label=" Click to load and test an example image"
        )
        
    # Event Handler
    classify_img_button.click(fn=process_img, inputs=upload_img, outputs=output_img)

if __name__ == "__main__":
    demo.launch()