Update app.py

app.py

@@ -1,43 +1,32 @@
 import torch
-import cv2
-import numpy as np
 import gradio as gr
+from torchvision.transforms import functional as F
 from PIL import Image
 
+# Load the YOLOv8 model (assuming it is already converted to the Hugging Face format)
+model = torch.hub.load('ultralytics/yolov8', 'custom', path='yolov5s.pt')
 
-model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
+# Define the prediction function
+def predict(image):
+    # Preprocess the input image
+    image_tensor = F.to_tensor(image)
+    image_tensor.unsqueeze_(0)
 
+    # Perform inference
+    results = model(image_tensor)
 
-model.conf = 0.25  
-model.iou = 0.45  
-model.agnostic = False  
-model.multi_label = False  
-model.max_det = 1000
+    # Post-process the results
+    # Extract the bounding box coordinates and class labels
+    bboxes = results.xyxy[0].tolist()
+    labels = results.names[0]
 
+    return bboxes, labels
 
-def detect(img):
+# Define the Gradio interface
+inputs = gr.inputs.Image()
+outputs = gr.outputs.Image()
 
+interface = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, capture_session=True)
 
-    results = model(img, size=640)
-
-    predictions = results.pred[0]
-    boxes = predictions[:, :4] # x1, y1, x2, y2
-    scores = predictions[:, 4]
-    categories = predictions[:, 5]
-    new_image = np.squeeze(results.render())
-
-    return new_image
-
-
-examples = ['apple_img.jpg', 'michael.gif']
-
-
-css = ".output-image, .input-image, .image-preview {height: 400px !important}"
-
-iface = gr.Interface(fn=detect, 
-                     inputs=gr.inputs.Image(type="numpy",), 
-                     outputs=gr.outputs.Image(type="numpy",),
-                     css=css,
-                     examples = examples,
-                     )
-iface.launch(debug=True, inline=True)
+# Run the interface
+interface.launch()