Update app.py
Browse files
app.py
CHANGED
@@ -4,8 +4,11 @@ from PIL import Image
|
|
4 |
from ultralyticsplus import YOLO, render_result
|
5 |
import cv2
|
6 |
import numpy as np
|
|
|
7 |
|
8 |
model = YOLO('best (1).pt')
|
|
|
|
|
9 |
|
10 |
# for i, r in enumerate(results):
|
11 |
|
@@ -17,6 +20,8 @@ def response(image):
|
|
17 |
print(image)
|
18 |
results = model(image)
|
19 |
text=""
|
|
|
|
|
20 |
for r in results:
|
21 |
conf = np.array(r.boxes.conf)
|
22 |
cls = np.array(r.boxes.cls)
|
@@ -30,13 +35,22 @@ def response(image):
|
|
30 |
conef = conef * 100
|
31 |
text += (f"Detected {name[cl]} with confidence {round(conef,1)}% at ({xy[0]},{xy[1]})\n")
|
32 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
33 |
|
34 |
# im_rgb = Image.fromarray(im_rgb)
|
35 |
|
36 |
|
37 |
-
return text
|
|
|
38 |
|
39 |
-
name = ['grenade','knife','pistol','rifle']
|
40 |
|
41 |
def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold: gr.Slider = 0.3, iou_threshold: gr.Slider = 0.6):
|
42 |
|
@@ -47,7 +61,7 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
|
|
47 |
render = render_result(model=model, image=image, result=results[0], rect_th = 1, text_th = 1)
|
48 |
|
49 |
|
50 |
-
|
51 |
|
52 |
|
53 |
# xywh = int(results.boxes.xywh)
|
@@ -56,7 +70,7 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
|
|
56 |
|
57 |
|
58 |
|
59 |
-
return render,
|
60 |
|
61 |
|
62 |
inputs = [
|
@@ -71,7 +85,8 @@ inputs = [
|
|
71 |
|
72 |
|
73 |
outputs = [gr.Image( type="filepath", label="Output Image"),
|
74 |
-
gr.Textbox(label="Result")
|
|
|
75 |
]
|
76 |
|
77 |
title = "YOLOv8 Custom Object Detection by Uyen Nguyen"
|
@@ -82,6 +97,7 @@ examples = [['th (11).jpg', 640, 0.3, 0.6],
|
|
82 |
['th (3).jpg', 640, 0.3, 0.6],
|
83 |
# ['four.jpg', 832, 0.3, 0.3]]
|
84 |
]
|
|
|
85 |
|
86 |
# yolo_app = gr.Interface(
|
87 |
# fn=yoloV8_func,
|
|
|
4 |
from ultralyticsplus import YOLO, render_result
|
5 |
import cv2
|
6 |
import numpy as np
|
7 |
+
from transformers import pipeline
|
8 |
|
9 |
model = YOLO('best (1).pt')
|
10 |
+
model2 = pipeline('Kaludi/csgo-weapon-classification')
|
11 |
+
name = ['grenade','knife','pistol','rifle']
|
12 |
|
13 |
# for i, r in enumerate(results):
|
14 |
|
|
|
20 |
print(image)
|
21 |
results = model(image)
|
22 |
text=""
|
23 |
+
name=""
|
24 |
+
|
25 |
for r in results:
|
26 |
conf = np.array(r.boxes.conf)
|
27 |
cls = np.array(r.boxes.cls)
|
|
|
35 |
conef = conef * 100
|
36 |
text += (f"Detected {name[cl]} with confidence {round(conef,1)}% at ({xy[0]},{xy[1]})\n")
|
37 |
|
38 |
+
if cl == 0:
|
39 |
+
name += name[cl]
|
40 |
+
elif cl == 1:
|
41 |
+
name += name[cl]
|
42 |
+
elif cl == 2:
|
43 |
+
name += model2(image)
|
44 |
+
elif cl == 3:
|
45 |
+
name += model2(image)
|
46 |
+
|
47 |
|
48 |
# im_rgb = Image.fromarray(im_rgb)
|
49 |
|
50 |
|
51 |
+
return name, text
|
52 |
+
|
53 |
|
|
|
54 |
|
55 |
def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold: gr.Slider = 0.3, iou_threshold: gr.Slider = 0.6):
|
56 |
|
|
|
61 |
render = render_result(model=model, image=image, result=results[0], rect_th = 1, text_th = 1)
|
62 |
|
63 |
|
64 |
+
weapon_name, text_detection = response(image)
|
65 |
|
66 |
|
67 |
# xywh = int(results.boxes.xywh)
|
|
|
70 |
|
71 |
|
72 |
|
73 |
+
return render, text_detection, weapon_name
|
74 |
|
75 |
|
76 |
inputs = [
|
|
|
85 |
|
86 |
|
87 |
outputs = [gr.Image( type="filepath", label="Output Image"),
|
88 |
+
gr.Textbox(label="Result"),
|
89 |
+
gr.Textbox(label="Weapon Name")
|
90 |
]
|
91 |
|
92 |
title = "YOLOv8 Custom Object Detection by Uyen Nguyen"
|
|
|
97 |
['th (3).jpg', 640, 0.3, 0.6],
|
98 |
# ['four.jpg', 832, 0.3, 0.3]]
|
99 |
]
|
100 |
+
|
101 |
|
102 |
# yolo_app = gr.Interface(
|
103 |
# fn=yoloV8_func,
|