import gradio as gr
import numpy as np
from PIL import Image
import subprocess
import sys

subprocess.check_call([sys.executable, "-m", "pip", "install", "ultralytics"])

from ultralytics import YOLO
import time
import os

# Reduce CPU overhead
os.environ["OMP_NUM_THREADS"] = "1"

MODEL_PATH = "best (1).pt"

# Load model once
model = YOLO(MODEL_PATH)

CLASS_COLORS = {
    "good": "green",
    "bad": "red"
}

def predict(image):
    if image is None:
        return "No image provided"

    # Convert PIL → numpy
    img = np.array(image)

    start = time.time()
    result = model(img, verbose=False)[0]
    latency = (time.time() - start) * 1000

    probs = result.probs.data.cpu().numpy()
    class_id = int(np.argmax(probs))
    confidence = float(probs[class_id]) * 100

    label = result.names[class_id]

    color = CLASS_COLORS.get(label, "black")

    output = (
        f"<h2 style='color:{color}; text-align:center;'>"
        f"{label.upper()}</h2>"
        f"<p style='text-align:center;'>"
        f"Confidence: <b>{confidence:.2f}%</b><br>"
        f"Inference Time: {latency:.1f} ms"
        f"</p>"
    )

    return output

with gr.Blocks() as demo:
    gr.Markdown("# 🔍 AI Simple Defect Classifier")
    gr.Markdown(
        "Upload or capture an image to classify an industrial part as **GOOD** or **BAD**."
    )

    with gr.Row():
        image_input = gr.Image(
            type="pil",
            label="Input Image",
            sources=["upload", "webcam"]
        )

    output = gr.HTML()

    classify_btn = gr.Button("Run Inspection")

    classify_btn.click(
        fn=predict,
        inputs=image_input,
        outputs=output
    )

demo.launch(theme=gr.themes.Soft(), share=True)