Abs6187 commited on
Commit
0d4b577
·
verified ·
1 Parent(s): e9f13f1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +16 -3
app.py CHANGED
@@ -1,6 +1,9 @@
1
  import gradio as gr
2
  from ultralytics import YOLO
3
  from PIL import Image
 
 
 
4
 
5
  # Load the trained YOLOv8 model
6
  model = YOLO("best.pt")
@@ -11,14 +14,24 @@ def predict(image):
11
  results_img = results[0].plot() # Get image with bounding boxes
12
  return Image.fromarray(results_img)
13
 
 
 
 
 
 
 
 
 
 
14
  # Create Gradio interface
15
  interface = gr.Interface(
16
  fn=predict,
17
  inputs=gr.Image(type="pil"),
18
  outputs=gr.Image(type="pil"),
19
  title="Helmet Detection with YOLOv8",
20
- description="Upload an image to detect helmets."
 
21
  )
22
 
23
- # Launch Gradio app
24
- interface.launch(share=True)
 
1
  import gradio as gr
2
  from ultralytics import YOLO
3
  from PIL import Image
4
+ import os
5
+
6
+
7
 
8
  # Load the trained YOLOv8 model
9
  model = YOLO("best.pt")
 
14
  results_img = results[0].plot() # Get image with bounding boxes
15
  return Image.fromarray(results_img)
16
 
17
+ # Get example images from the images folder
18
+ def get_example_images():
19
+ examples = []
20
+ image_folder = "images"
21
+ for filename in os.listdir(image_folder):
22
+ if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
23
+ examples.append(os.path.join(image_folder, filename))
24
+ return examples
25
+
26
  # Create Gradio interface
27
  interface = gr.Interface(
28
  fn=predict,
29
  inputs=gr.Image(type="pil"),
30
  outputs=gr.Image(type="pil"),
31
  title="Helmet Detection with YOLOv8",
32
+ description="Upload an image to detect helmets.",
33
+ examples=get_example_images()
34
  )
35
 
36
+ # Launch the interface
37
+ interface.launch(share=True)