import os
from collections import Counter

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

# ==========================================================
# Load Model
# ==========================================================

model = YOLO("best.pt")

# ==========================================================
# Example Images
# ==========================================================

example_images = []

if os.path.exists("examples"):
    for file in sorted(os.listdir("examples")):
        if file.lower().endswith((".jpg", ".jpeg", ".png")):
            example_images.append([os.path.join("examples", file)])

# ==========================================================
# Detection Function
# ==========================================================

def detect(image, conf, iou):

    results = model.predict(
        source=image,
        conf=conf,
        iou=iou,
        verbose=False
    )

    result = results[0]

    plotted = result.plot()

    plotted = cv2.cvtColor(plotted, cv2.COLOR_BGR2RGB)

    detected = []

    for cls in result.boxes.cls.tolist():
        detected.append(model.names[int(cls)])

    counter = Counter(detected)

    table = []

    for name, count in sorted(counter.items()):
        table.append([name, count])

    return Image.fromarray(plotted), table

# ==========================================================
# Metric Images
# ==========================================================

metric_files = [
    "metrics/results.png",
    "metrics/P_curve.png",
    "metrics/R_curve.png",
    "metrics/PR_curve.png",
    "metrics/F1_curve.png",
    "metrics/confusion_matrix.png"
]

metric_components = []

for path in metric_files:
    if os.path.exists(path):
        metric_components.append(gr.Image(value=path, label=os.path.basename(path)))

# ==========================================================
# About Text
# ==========================================================

about = """
# Warehouse Vision AI

### Industrial Object Detection using YOLOv8

This project detects warehouse objects using a custom-trained YOLO model.

### Features

- Industrial Rack Detection
- KLT Box Detection
- Computer Hardware Detection
- Safety Equipment Detection
- Warehouse Object Localization

### Framework

- Ultralytics YOLO
- Gradio
- Hugging Face Spaces

### Author

Omkar Kalburgi
"""

# ==========================================================
# UI
# ==========================================================

with gr.Blocks(title="Warehouse Vision AI") as demo:

    gr.Markdown(
        """
# 📦 Warehouse Vision AI

### YOLO-based Industrial Warehouse Object Detection

Upload an image or try one of the sample images.
"""
    )

    with gr.Tabs():

        # --------------------------------------------------

        with gr.Tab("🔍 Detection"):

            with gr.Row():

                with gr.Column():

                    image = gr.Image(type="pil", label="Input Image")

                    conf = gr.Slider(
                        0.1,
                        1.0,
                        value=0.25,
                        step=0.05,
                        label="Confidence Threshold",
                    )

                    iou = gr.Slider(
                        0.1,
                        1.0,
                        value=0.45,
                        step=0.05,
                        label="IoU Threshold",
                    )

                    btn = gr.Button("Run Detection")

                with gr.Column():

                    output = gr.Image(label="Prediction")

                    table = gr.Dataframe(
                        headers=["Class", "Count"],
                        datatype=["str", "number"],
                        interactive=False,
                        label="Detected Objects",
                    )

            btn.click(
                detect,
                inputs=[image, conf, iou],
                outputs=[output, table],
            )

        # --------------------------------------------------

        with gr.Tab("🧪 Sample Images"):

            gr.Markdown("Click any image below to test the model.")

            sample_input = gr.Image(type="pil")

            sample_output = gr.Image()

            sample_table = gr.Dataframe(
                headers=["Class", "Count"],
                interactive=False,
            )

            gr.Examples(
                examples=example_images,
                inputs=sample_input,
            )

            sample_btn = gr.Button("Run Detection")

            sample_btn.click(
                detect,
                inputs=[sample_input, conf, iou],
                outputs=[sample_output, sample_table],
            )

        # --------------------------------------------------

        with gr.Tab("📈 Model Performance"):

            gr.Markdown("Training Metrics")

            for img in metric_components:
                img.render()

        # --------------------------------------------------

        with gr.Tab("📖 About"):

            gr.Markdown(about)

demo.launch()