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

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +43 -43
app.py CHANGED
@@ -1,43 +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()
 
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 = '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()