itsanurag commited on
Commit
a52905f
·
verified ·
1 Parent(s): a1ed03d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +69 -0
app.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from ultralyticsplus import YOLO, render_result
4
+
5
+
6
+ def yoloV9_func(image: gr.inputs.Image = None,
7
+ image_size: gr.inputs.Slider = 640,
8
+ conf_threshold: gr.inputs.Slider = 0.4,
9
+ iou_threshold: gr.inputs.Slider = 0.50):
10
+ """This function performs YOLOv9 object detection on the given image.
11
+
12
+ Args:
13
+ image (gr.inputs.Image, optional): Input image to detect objects on. Defaults to None.
14
+ image_size (gr.inputs.Slider, optional): Desired image size for the model. Defaults to 640.
15
+ conf_threshold (gr.inputs.Slider, optional): Confidence threshold for object detection. Defaults to 0.4.
16
+ iou_threshold (gr.inputs.Slider, optional): Intersection over Union threshold for object detection. Defaults to 0.50.
17
+ """
18
+ # Load the YOLOv8 model from the 'best.pt' checkpoint
19
+ model_path = "best_model.pt"
20
+ model = YOLO(model_path)
21
+
22
+ # Perform object detection on the input image using the YOLOv8 model
23
+ results = model.predict(image,
24
+ conf=conf_threshold,
25
+ iou=iou_threshold,
26
+ imgsz=image_size)
27
+
28
+ # Print the detected objects' information (class, coordinates, and probability)
29
+ box = results[0].boxes
30
+ print("Object type:", box.cls)
31
+ print("Coordinates:", box.xyxy)
32
+ print("Probability:", box.conf)
33
+
34
+ # Render the output image with bounding boxes around detected objects
35
+ render = render_result(model=model, image=image, result=results[0])
36
+ return render
37
+
38
+
39
+ inputs = [
40
+ gr.inputs.Image(type="filepath", label="Input Image"),
41
+ gr.inputs.Slider(minimum=320, maximum=1280, default=640,
42
+ step=32, label="Image Size"),
43
+ gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25,
44
+ step=0.05, label="Confidence Threshold"),
45
+ gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45,
46
+ step=0.05, label="IOU Threshold"),
47
+ ]
48
+
49
+
50
+ outputs = gr.outputs.Image(type="filepath", label="Output Image")
51
+
52
+ title = "CUSTOM yolov9 model for room cleanliness"
53
+
54
+
55
+ examples = [['one.jpg', 640, 0.5, 0.7],
56
+ ['two.jpg', 640, 0.5, 0.6],
57
+ ['three.jpg', 640, 0.5, 0.8]]
58
+
59
+ yolo_app = gr.Interface(
60
+ fn=yoloV9_func,
61
+ inputs=inputs,
62
+ outputs=outputs,
63
+ title=title,
64
+ examples=examples,
65
+ cache_examples=True,
66
+ )
67
+
68
+ # Launch the Gradio interface in debug mode with queue enabled
69
+ yolo_app.launch(debug=True, enable_queue=True)