Update app.py

app.py

@@ -1,40 +1,26 @@
 import gradio as gr
 from ultralytics import YOLO
-from PIL import Image
-import numpy as np
 import os
 
-# Load model (make sure best.pt is in the same directory)
 MODEL_PATH = "best.pt"
 
 if not os.path.exists(MODEL_PATH):
-    raise FileNotFoundError("best.pt not found. Upload your trained model to the Space.")
+    raise FileNotFoundError("best.pt not found.")
 
 model = YOLO(MODEL_PATH)
 
 def predict(image):
-    """
-    Run YOLOv8 pose inference on uploaded image
-    """
     results = model(image)
+    return results[0].plot()
 
-    # Plot result with keypoints
-    annotated_frame = results[0].plot()
+with gr.Blocks() as demo:
+    gr.Markdown("# 🐶 DogFLW Pose Detection")
+    gr.Markdown("Upload an image to detect dog facial landmarks.")
 
-    return annotated_frame
+    input_img = gr.Image(type="pil")
+    output_img = gr.Image()
 
-title = "🐶 DogFLW YOLOv8 Pose Detection"
-description = """
-Upload an image of a dog to detect facial keypoints using a trained YOLOv8 pose model.
-"""
+    btn = gr.Button("Run Detection")
+    btn.click(predict, inputs=input_img, outputs=output_img)
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="numpy"),
-    title=title,
-    description=description,
-)
-
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()