Exched commited on
Commit
5d75401
·
verified ·
1 Parent(s): 074e85f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -0
app.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from PIL import Image
4
+ import requests
5
+ from io import BytesIO
6
+
7
+ # Load YOLOv5 pre-trained model from Hugging Face
8
+ model = torch.hub.load('ultralytics/yolov5', 'yolov5s') # You can choose other versions like yolov5m or yolov5l
9
+
10
+ # Function for object detection
11
+ def detect_objects(input_image):
12
+ # If the input is a URL, download the image
13
+ if isinstance(input_image, str):
14
+ response = requests.get(input_image)
15
+ img = Image.open(BytesIO(response.content))
16
+ else:
17
+ img = Image.fromarray(input_image)
18
+
19
+ # Run YOLOv5 object detection
20
+ results = model(img)
21
+
22
+ # Render results on image
23
+ results.render() # Render boxes on the image
24
+
25
+ # Return image with detections
26
+ output_image = results.imgs[0]
27
+ return Image.fromarray(output_image)
28
+
29
+ # Create Gradio interface
30
+ interface = gr.Interface(
31
+ fn=detect_objects,
32
+ inputs=gr.inputs.Image(type="numpy", label="Upload an image"),
33
+ outputs=gr.outputs.Image(type="pil", label="Detected Image"),
34
+ title="YOLOv5 Object Detection",
35
+ description="Upload an image and detect objects using YOLOv5 model. The model can identify objects like people, cars, animals, and more.",
36
+ theme="huggingface"
37
+ )
38
+
39
+ # Launch the interface
40
+ interface.launch()