zliang commited on
Commit
74f482e
1 Parent(s): e304110

Upload app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -0
app.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from ultralytics import YOLO
3
+ import cv2
4
+ import numpy as np
5
+
6
+ # Load the trained model
7
+ model_path = 'runs\\detect\\train6\\weights\\best.pt' # Replace with the path to your trained .pt file
8
+ model = YOLO(model_path)
9
+
10
+ # Function to perform inference on an image
11
+ def infer_image(image):
12
+ # Convert the image from BGR to RGB
13
+ image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
14
+
15
+ # Perform inference
16
+ results = model(image_rgb)
17
+
18
+ # Extract results and annotate image
19
+ for result in results:
20
+ for box in result.boxes:
21
+ x1, y1, x2, y2 = box.xyxy[0]
22
+ cls = int(box.cls[0])
23
+ conf = float(box.conf[0])
24
+
25
+ # Draw bounding box
26
+ cv2.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
27
+ # Draw label
28
+ label = f'{model.names[cls]} {conf:.2f}'
29
+ cv2.putText(image, label, (int(x1), int(y1) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
30
+
31
+ return image
32
+
33
+ # Create Gradio interface
34
+ iface = gr.Interface(
35
+ fn=infer_image,
36
+ inputs=gr.Image(type="numpy", label="Upload an Image"),
37
+ outputs=gr.Image(type="numpy", label="Annotated Image"),
38
+ title="YOLOv8 Inference",
39
+ description="Upload an image to get object detection results using YOLOv8."
40
+ )
41
+
42
+ # Launch the app
43
+ iface.launch()