scchess commited on
Commit
7c99a6d
·
1 Parent(s): 2c5b705

Added detection boxes

Browse files
Files changed (1) hide show
  1. app.py +21 -3
app.py CHANGED
@@ -9,14 +9,32 @@ model = YOLO('best.pt') # Make sure best.pt is in the same folder
9
  def detect(image):
10
  results = model(image)
11
  annotated_frame = results[0].plot() # Draw predictions on the image
12
- return Image.fromarray(annotated_frame)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
13
 
14
  # Launch the Gradio interface
15
  gr.Interface(
16
  fn=detect,
17
  inputs=gr.Image(type="pil"),
18
- outputs=gr.Image(type="pil"),
 
 
 
19
  title="YOLOv9 (Ultralytics) Object Detection",
20
- description="Upload an image to run object detection using your custom YOLOv9 model.",
21
  ).launch()
22
 
 
9
  def detect(image):
10
  results = model(image)
11
  annotated_frame = results[0].plot() # Draw predictions on the image
12
+
13
+ # Get detection results
14
+ detections = []
15
+ for r in results:
16
+ boxes = r.boxes
17
+ for box in boxes:
18
+ x1, y1, x2, y2 = box.xyxy[0].tolist() # Get box coordinates
19
+ conf = float(box.conf[0]) # Get confidence
20
+ cls = int(box.cls[0]) # Get class
21
+ detections.append({
22
+ 'class': cls,
23
+ 'confidence': conf,
24
+ 'box': [x1, y1, x2, y2]
25
+ })
26
+
27
+ return Image.fromarray(annotated_frame), detections
28
 
29
  # Launch the Gradio interface
30
  gr.Interface(
31
  fn=detect,
32
  inputs=gr.Image(type="pil"),
33
+ outputs=[
34
+ gr.Image(type="pil", label="Detected Image"),
35
+ gr.JSON(label="Detection Results")
36
+ ],
37
  title="YOLOv9 (Ultralytics) Object Detection",
38
+ description="Upload an image to run object detection using your custom YOLOv9 model. The results will show both the annotated image and the detection details.",
39
  ).launch()
40