import gradio as gr
print("[STARTUP] Imports completed (gradio, ultralytics, PIL loaded)", flush=True)

from ultralytics import YOLO
from PIL import Image

print("[STARTUP] Loading YOLO model from best.pt ...", flush=True)
model = YOLO("best.pt")
print("[STARTUP] Model loaded successfully.", flush=True)
print(f"[STARTUP] Model classes: {model.names}", flush=True)


def detect_damage(input_image: Image.Image):
    print("[REQUEST] Received new image for detection", flush=True)

    if input_image is None:
        return None, "No image provided."

    results = model.predict(
        source=input_image,
        imgsz=640,
        conf=0.4,
        iou=0.45,
    )

    result = results[0]
    result_image = Image.fromarray(result.plot()[:, :, ::-1])

    detections = result.boxes
    if detections is None or len(detections) == 0:
        summary = "No damage detected."
    else:
        lines = [f"Found {len(detections)} detection(s):"]
        for box in detections:
            class_id = int(box.cls[0])
            confidence = float(box.conf[0])
            class_name = model.names[class_id]
            lines.append(f"- {class_name} (confidence: {confidence:.2f})")
        summary = "\n".join(lines)

    return result_image, summary


demo = gr.Interface(
    fn=detect_damage,
    inputs=gr.Image(type="pil", label="Upload an image"),
    outputs=[
        gr.Image(type="pil", label="Detection result"),
        gr.Textbox(label="Summary"),
    ],
    title="Structural Damage Detection (Crack / Spalling / Pothole)",
    description="Upload an image of a road or concrete structure to detect cracks, spalling, and potholes using a YOLOv8 model.",
)

print("[STARTUP] Calling demo.launch() now ...", flush=True)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)