app.py

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

# Load model (you can switch to 'fasterrcnn' based loading if needed)
model = YOLO("model/best.pt")  # fine-tuned YOLOv8 model

# Define defect labels if custom
DEFECT_LABELS = ['crack', 'spalling', 'rust', 'deformation']

# Inference function
def detect_defects(image):
    results = model(image)
    annotated_img = results[0].plot()  # Draw boxes
    predictions = results[0].boxes.data.cpu().numpy()  # [x1, y1, x2, y2, conf, class]

    # Create readable output
    output = []
    for pred in predictions:
        x1, y1, x2, y2, conf, cls_id = pred
        label = DEFECT_LABELS[int(cls_id)]
        output.append(f"{label}: {conf:.2f}")

    return Image.fromarray(annotated_img), "\n".join(output)

# Gradio UI
interface = gr.Interface(
    fn=detect_defects,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil", label="Detected Image"), gr.Textbox(label="Detected Faults")],
    title="Structural Defect Detection",
    description="Upload drone-captured image to detect cracks, rust, spalling, and deformations.",
    allow_flagging="never"
)

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