import gradio as gr
from ultralytics import YOLO
import cv2

model = YOLO('best.pt')

def predict_terrain(image):
    if image is None: return None
    results = model(image, imgsz=800)
    return cv2.cvtColor(results[0].plot(), cv2.COLOR_BGR2RGB)

# FIXED: Changed 'Slate()' to 'Soft()'
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("<h1 style='text-align: center;'>Pathfinder AI: Off-Road Terrain Segmentation</h1>")
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="numpy", label="Raw Input")
            btn = gr.Button("Initialize Segmentation", variant="primary")
        with gr.Column():
            img_output = gr.Image(type="numpy", label="Predicted Output")
    btn.click(fn=predict_terrain, inputs=img_input, outputs=img_output)

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