cdcvd commited on
Commit
92e9e16
·
verified ·
1 Parent(s): cb06d05

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +42 -0
app.py ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from PIL import Image
3
+ import numpy as np
4
+ import cv2
5
+ from ultralytics import YOLO
6
+
7
+ # Load the YOLOv8 model
8
+ model = YOLO('best.pt')
9
+
10
+ def detect_objects(image):
11
+ # Convert the input image to a format YOLO can work with
12
+ image = np.array(image)
13
+
14
+ # Perform detection
15
+ results = model(image)[0]
16
+
17
+ # Draw bounding boxes on the image
18
+ for box in results.boxes.data.cpu().numpy():
19
+ x1, y1, x2, y2, score, class_id = box
20
+ x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
21
+
22
+ # Draw the bounding box
23
+ cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
24
+
25
+ # Put the class name above the bounding box
26
+ class_name = model.model.names[int(class_id)]
27
+ cv2.putText(image, f'{class_name} {score:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
28
+
29
+ # Convert back to PIL image
30
+ return Image.fromarray(image)
31
+
32
+ # Define the Gradio interface
33
+ interface = gr.Interface(
34
+ fn=detect_objects,
35
+ inputs=gr.inputs.Image(type="pil"),
36
+ outputs=gr.outputs.Image(type="pil"),
37
+ title="YOLOv8 Object Detection",
38
+ description="Upload an image and YOLOv8 will detect objects in the image."
39
+ )
40
+
41
+ # Launch the interface
42
+ interface.launch()