Varun commited on
Commit
63bef64
·
1 Parent(s): 644c203

Refactor object detection in app.py to use YOLOv8 model from Hugging Face, removing DiffusionPipeline. Update requirements.txt to include ultralytics and numpy, and adjust Gradio output type.

Browse files
Files changed (2) hide show
  1. app.py +10 -23
  2. requirements.txt +3 -3
app.py CHANGED
@@ -1,37 +1,24 @@
1
- import os
2
  import gradio as gr
3
- from diffusers import DiffusionPipeline
4
- import torch
5
- from huggingface_hub import login
6
 
7
-
8
- login(token=os.getenv("HF_TOKEN"))
9
-
10
- # Initialize the pipeline
11
- pipe = DiffusionPipeline.from_pretrained("Lookingsoft-team/object_detection")
12
- if torch.cuda.is_available():
13
- pipe = pipe.to("cuda")
14
 
15
 
16
  def detect_objects(image):
17
- if image is None:
18
- return None
19
-
20
- # Process the image through the pipeline
21
- # Note: This is a placeholder - actual processing will depend on the model's specific requirements
22
- results = pipe(image=image)
23
-
24
- # Return the processed image with detections
25
- return results.images[0]
26
 
27
 
28
- # Create Gradio interface
29
  demo = gr.Interface(
30
  fn=detect_objects,
31
  inputs=gr.Image(type="pil"),
32
- outputs=gr.Image(type="pil"),
33
  title="Object Detection",
34
- description="Upload an image and the model will detect objects in it!",
35
  )
36
 
37
  if __name__ == "__main__":
 
 
1
  import gradio as gr
2
+ from ultralytics import YOLO
 
 
3
 
4
+ # Load YOLOv8 model from Hugging Face
5
+ model = YOLO(
6
+ "https://huggingface.co/Lookingsoft-team/object_detection/resolve/main/yolov8n.pt"
7
+ )
 
 
 
8
 
9
 
10
  def detect_objects(image):
11
+ results = model(image)
12
+ annotated_image = results[0].plot() # Draw bounding boxes
13
+ return annotated_image
 
 
 
 
 
 
14
 
15
 
 
16
  demo = gr.Interface(
17
  fn=detect_objects,
18
  inputs=gr.Image(type="pil"),
19
+ outputs=gr.Image(type="numpy"),
20
  title="Object Detection",
21
+ description="Upload an image for object detection!",
22
  )
23
 
24
  if __name__ == "__main__":
requirements.txt CHANGED
@@ -1,5 +1,5 @@
1
  gradio>=5.32.1
 
2
  torch
3
- diffusers
4
- transformers
5
- accelerate
 
1
  gradio>=5.32.1
2
+ ultralytics
3
  torch
4
+ numpy
5
+ Pillow