Dricz commited on
Commit
80fb3d6
1 Parent(s): 599b5b5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +30 -41
app.py CHANGED
@@ -40,47 +40,13 @@ video_directory = "/home/user/app/video"
40
  # im_bgr = r.plot()
41
  # im_rgb = im_bgr[..., ::-1] # Convert BGR to RGB
42
 
43
- def response(image):
44
- print(image)
45
- results = model(image)
46
- text = ""
47
- name_weap = ""
48
-
49
- for r in results:
50
- conf = np.array(r.boxes.conf)
51
- cls = np.array(r.boxes.cls)
52
- cls = cls.astype(int)
53
- xywh = np.array(r.boxes.xywh)
54
- xywh = xywh.astype(int)
55
-
56
- for con, cl, xy in zip(conf, cls, xywh):
57
- cone = con.astype(float)
58
- conef = round(cone,3)
59
- conef = conef * 100
60
- text += (f"Detected {name[cl]} with confidence {round(conef,1)}% at ({xy[0]},{xy[1]})\n")
61
-
62
- if cl == 0:
63
- name_weap += name[cl] + '\n'
64
- elif cl == 1:
65
- name_weap += name[cl] + '\n'
66
- elif cl == 2:
67
- out = model2(image)
68
- name_weap += out[0]["label"] + '\n'
69
- elif cl == 3:
70
- out = model2(image)
71
- name_weap += out[0]["label"] + '\n'
72
-
73
-
74
- # im_rgb = Image.fromarray(im_rgb)
75
-
76
-
77
- return name_weap, text
78
-
79
-
80
 
81
  def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold: gr.Slider = 0.3, iou_threshold: gr.Slider = 0.6):
82
 
83
- results = model.predict(image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size)
 
 
 
84
 
85
  box = results[0].boxes
86
 
@@ -89,14 +55,37 @@ def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold
89
  im = Image.fromarray(im_array[..., ::-1])
90
 
91
 
92
- weapon_name, text_detection = response(image)
93
-
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
94
 
95
  # xywh = int(results.boxes.xywh)
96
  # x = xywh[0]
97
  # y = xywh[1]
98
 
99
- return im, text_detection, weapon_name
100
 
101
 
102
  inputs = [
 
40
  # im_bgr = r.plot()
41
  # im_rgb = im_bgr[..., ::-1] # Convert BGR to RGB
42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
43
 
44
  def response2(image: gr.Image = None,image_size: gr.Slider = 640, conf_threshold: gr.Slider = 0.3, iou_threshold: gr.Slider = 0.6):
45
 
46
+ results = model.predict(image, conf=conf_threshold, iou=iou_threshold, imgsz=image_size
47
+
48
+ text = ""
49
+ name_weap = ""
50
 
51
  box = results[0].boxes
52
 
 
55
  im = Image.fromarray(im_array[..., ::-1])
56
 
57
 
58
+
59
+ for r in results:
60
+ conf = np.array(r.boxes.conf)
61
+ cls = np.array(r.boxes.cls)
62
+ cls = cls.astype(int)
63
+ xywh = np.array(r.boxes.xywh)
64
+ xywh = xywh.astype(int)
65
+
66
+ for con, cl, xy in zip(conf, cls, xywh):
67
+ cone = con.astype(float)
68
+ conef = round(cone,3)
69
+ conef = conef * 100
70
+ text += (f"Detected {name[cl]} with confidence {round(conef,1)}% at ({xy[0]},{xy[1]})\n")
71
+
72
+ if cl == 0:
73
+ name_weap += name[cl] + '\n'
74
+ elif cl == 1:
75
+ name_weap += name[cl] + '\n'
76
+ elif cl == 2:
77
+ out = model2(image)
78
+ name_weap += out[0]["label"] + '\n'
79
+ elif cl == 3:
80
+ out = model2(image)
81
+ name_weap += out[0]["label"] + '\n'
82
+
83
 
84
  # xywh = int(results.boxes.xywh)
85
  # x = xywh[0]
86
  # y = xywh[1]
87
 
88
+ return im, text, name_weap
89
 
90
 
91
  inputs = [