syedfaisalabrar commited on
Commit
f0c1530
·
verified ·
1 Parent(s): c0e5682

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -4
app.py CHANGED
@@ -6,25 +6,35 @@ from PIL import Image
6
  from ultralytics import YOLO
7
 
8
  # Load YOLOv11 Model
9
- model_path = "best.pt"
10
  model = YOLO(model_path)
11
 
12
  def predict(image):
13
  image = np.array(image)
14
  results = model(image)
15
 
16
- # Draw bounding boxes
 
17
  for result in results:
18
  for box in result.boxes:
19
  x1, y1, x2, y2 = map(int, box.xyxy[0])
20
  conf = box.conf[0]
21
  cls = int(box.cls[0])
22
  label = f"{model.names[cls]} {conf:.2f}"
 
 
 
23
  cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
24
  cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
25
 
26
- return Image.fromarray(image)
27
 
28
  # Gradio Interface
29
- iface = gr.Interface(fn=predict, inputs="image", outputs="image", title="YOLOv11 Object Detection")
 
 
 
 
 
 
30
  iface.launch()
 
6
  from ultralytics import YOLO
7
 
8
  # Load YOLOv11 Model
9
+ model_path = "best.pt"
10
  model = YOLO(model_path)
11
 
12
  def predict(image):
13
  image = np.array(image)
14
  results = model(image)
15
 
16
+ labels = []
17
+ # Draw bounding boxes and extract labels
18
  for result in results:
19
  for box in result.boxes:
20
  x1, y1, x2, y2 = map(int, box.xyxy[0])
21
  conf = box.conf[0]
22
  cls = int(box.cls[0])
23
  label = f"{model.names[cls]} {conf:.2f}"
24
+
25
+ labels.append(label) # Store detected labels
26
+
27
  cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
28
  cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
29
 
30
+ return Image.fromarray(image), labels
31
 
32
  # Gradio Interface
33
+ iface = gr.Interface(
34
+ fn=predict,
35
+ inputs="image",
36
+ outputs=["image", "text"], # Returning both image and detected labels
37
+ title="YOLOv11 Object Detection"
38
+ )
39
+
40
  iface.launch()