Pathfinder-AI / app.py
buoipomn's picture
Update app.py
c4a580a verified
Raw
History Blame Contribute Delete
837 Bytes
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()