@@ -72,6 +72,8 @@ def draw_box_with_text(img, rect_list, score_list):
72
72
def nms (rect_list , score_list ):
73
73
"""
74
74
非最大抑制
75
+ :param rect_list: list,大小为[N, 4]
76
+ :param score_list: list,大小为[N]
75
77
"""
76
78
nms_rects = list ()
77
79
nms_scores = list ()
@@ -141,6 +143,9 @@ def nms(rect_list, score_list):
141
143
# 保存正样本边界框以及
142
144
score_list = list ()
143
145
positive_list = list ()
146
+
147
+ # tmp_score_list = list()
148
+ # tmp_positive_list = list()
144
149
for rect in rects :
145
150
xmin , ymin , xmax , ymax = rect
146
151
rect_img = img [ymin :ymax , xmin :xmax ]
@@ -154,12 +159,23 @@ def nms(rect_list, score_list):
154
159
"""
155
160
probs = torch .softmax (output , dim = 0 ).cpu ().numpy ()
156
161
162
+ # tmp_score_list.append(probs[1])
163
+ # tmp_positive_list.append(rect)
164
+
157
165
if probs [1 ] >= svm_thresh :
158
166
score_list .append (probs [1 ])
159
167
positive_list .append (rect )
160
168
# cv2.rectangle(dst, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2)
161
169
print (rect , output , probs )
162
170
171
+ # tmp_img2 = copy.deepcopy(dst)
172
+ # draw_box_with_text(tmp_img2, tmp_positive_list, tmp_score_list)
173
+ # cv2.imshow('tmp', tmp_img2)
174
+ #
175
+ # tmp_img = copy.deepcopy(dst)
176
+ # draw_box_with_text(tmp_img, positive_list, score_list)
177
+ # cv2.imshow('tmp2', tmp_img)
178
+
163
179
nms_rects , nms_scores = nms (positive_list , score_list )
164
180
print (nms_rects )
165
181
print (nms_scores )
0 commit comments