Update app.py

app.py

@@ -1,14 +1,39 @@
+
+
+import torch
+import cv2
 import gradio as gr
+from pathlib import Path
+
+# Load the YOLOv5 model
+MODEL_PATH = "best.pt"  # Path to your custom model
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model = torch.hub.load('ultralytics/yolov5', 'custom', path=MODEL_PATH, force_reload=True).to(device)
+
+# Define the prediction function
+def predict(image):
+    # Convert the input image (numpy) to a compatible format
+    img_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+    results = model(img_bgr)  # Run inference
 
-from ultralytics import YOLO
+    # Render results
+    annotated_image = results.render()[0]
+    annotated_image_rgb = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
+    return annotated_image_rgb
 
-model = YOLO("best.pt")
+# Gradio Interface
+title = "Custom YOLOv5 Model Deployment"
+description = "Upload an image to get predictions using a custom YOLOv5 model trained with your dataset."
 
-def greet(source):
-  model = YOLO("best.pt")
-  results = model.predict(source, imgsz=1080, conf=0.25)
-  result = results[0]
-  return result.plot()
+interface = gr.Interface(
+    fn=predict,                  # Prediction function
+    inputs=gr.Image(type="numpy", label="Upload Image"),
+    outputs=gr.Image(type="numpy", label="Detected Objects"),
+    title=title,
+    description=description,
+    live=True  # Enable live processing
+)
 
-demo = gr.Interface(fn=greet, inputs="image", outputs="image")
-demo.launch()
+# Launch the Gradio app
+if __name__ == "__main__":
+    interface.launch()