2lu commited on
Commit
41e7e8e
1 Parent(s): f966204

model upload

Browse files
Files changed (4) hide show
  1. app.py +145 -0
  2. model/yolo_efficient.onnx +3 -0
  3. requirements.txt +2 -0
  4. samples/out_1.jpg +0 -0
app.py ADDED
@@ -0,0 +1,145 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import onnxruntime as rt
3
+ import cv2
4
+ import numpy as np
5
+ from PIL import Image
6
+
7
+ H, W = 224, 224
8
+ classes=['aeroplane','bicycle','bird','boat','bottle','bus','car','cat','chair','cow','diningtable',
9
+ 'dog','horse','motorbike','person','pottedplant','sheep','sofa','train','tvmonitor']
10
+
11
+ providers = ['CPUExecutionProvider']
12
+
13
+ m = rt.InferenceSession("./model/yolo_efficient.onnx", providers=providers)
14
+
15
+ def nms(final_boxes, scores, IOU_threshold=0):
16
+ scores = np.array(scores)
17
+ final_boxes = np.array(final_boxes)
18
+
19
+ boxes = final_boxes[...,:-1]
20
+
21
+ boxes = [list(map(int, i)) for i in boxes]
22
+ boxes = np.array(boxes)
23
+ #print(boxes)
24
+
25
+ x1 = boxes[:, 0]
26
+ y1 = boxes[:, 1]
27
+ x2 = boxes[:, 2]
28
+ y2 = boxes[:, 3]
29
+
30
+
31
+ area = (x2 - x1)*(y2 - y1)
32
+
33
+ order = np.argsort(scores)
34
+ #print(order)
35
+
36
+ pick = []
37
+
38
+ while len(order) > 0:
39
+ last = len(order)-1
40
+ i = order[last]
41
+ pick.append(i)
42
+
43
+ suppress = [last]
44
+
45
+ if len(order)==0:
46
+ break
47
+
48
+ for pos in range(last):
49
+ j = order[pos]
50
+
51
+ xx1 = max(x1[i], x1[j])
52
+ yy1 = max(y1[i], y1[j])
53
+ xx2 = min(x2[i], x2[j])
54
+ yy2 = min(y2[i], y2[j])
55
+
56
+ w = max(0, xx2-xx1+1)
57
+ h = max(0, yy2-yy1+1)
58
+
59
+ overlap = float(w*h)/area[j]
60
+
61
+ if overlap > IOU_threshold:
62
+ suppress.append(pos)
63
+
64
+ order = np.delete(order, suppress)
65
+
66
+ return final_boxes[pick]
67
+
68
+ def detect_obj(input_image):
69
+ try:
70
+ image = np.array(input_image)
71
+ image = cv2.resize(image, (H, W))
72
+ img = image
73
+
74
+ image = image.astype(np.float32)
75
+ image = np.expand_dims(image, axis=0)
76
+ print(image.shape)
77
+
78
+ output = m.run(['reshape'], {"input": image})
79
+ output = np.squeeze(output, axis=0)
80
+ print(output.shape)
81
+
82
+ THRESH=.25
83
+
84
+
85
+ object_positions = np.concatenate(
86
+ [np.stack(np.where(output[..., 0]>=THRESH), axis=-1),
87
+ np.stack(np.where(output[..., 5]>=THRESH), axis=-1)], axis=0
88
+ )
89
+
90
+ selected_output = []
91
+ for indices in object_positions:
92
+ selected_output.append(output[indices[0]][indices[1]][indices[2]])
93
+ selected_output = np.array(selected_output)
94
+
95
+ final_boxes = []
96
+ final_scores = []
97
+
98
+ for i,pos in enumerate(object_positions):
99
+ for j in range(2):
100
+ if selected_output[i][j*5]>THRESH:
101
+ output_box = np.array(output[pos[0]][pos[1]][pos[2]][(j*5)+1:(j*5)+5], dtype=float)
102
+
103
+ x_centre = (np.array(pos[1], dtype=float) + output_box[0])*32
104
+ y_centre = (np.array(pos[2], dtype=float) + output_box[1])*32
105
+
106
+ x_width, y_height = abs(W*output_box[2]), abs(H*output_box[3])
107
+
108
+ x_min, y_min = int(x_centre - (x_width/2)), int(y_centre-(y_height/2))
109
+ x_max, y_max = int(x_centre+(x_width/2)), int(y_centre + (y_height/2))
110
+
111
+ if(x_min<0):x_min=0
112
+ if(y_min<0):y_min=0
113
+ if(x_max<0):x_max=0
114
+ if(y_max<0):y_max=0
115
+
116
+ final_boxes.append(
117
+ [x_min, y_min, x_max, y_max, str(classes[np.argmax(selected_output[..., 10:], axis=-1)[i]])]
118
+ )
119
+ final_scores.append(selected_output[i][j*5])
120
+
121
+ final_boxes = np.array(final_boxes)
122
+
123
+ nms_output = nms(final_boxes, final_scores, 0.3)
124
+
125
+ print(nms_output)
126
+ for i in nms_output:
127
+ cv2.rectangle(
128
+ img,
129
+ (int(i[0]), int(i[1])),
130
+ (int(i[2]), int(i[3])), (255, 0, 0)
131
+ )
132
+
133
+ cv2.putText(
134
+ img,
135
+ i[-1],
136
+ (int(i[0]), int(i[1])+15),
137
+ cv2.FONT_HERSHEY_PLAIN, 1, (255, 0, 0), 1
138
+ )
139
+
140
+ output_pil_img = Image.fromarray(np.uint8(img)).convert('RGB')
141
+
142
+ return output_pil_img
143
+
144
+ except:
145
+ return input_image
model/yolo_efficient.onnx ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:94649900941ab5e8e91d9446ca9ee8d5f3c974f13b7357bd8ded5a297ac797b3
3
+ size 132372993
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ onnxruntime
2
+ opencv-python-headless
samples/out_1.jpg ADDED