def get_top_cls(self, pred_bbox, org_h, org_w, top_n):
'''
获取top_n，类别和得分
:param pred_bbox:所有框
:param org_h:高
:param org_w:宽
:param top_n:top数
:return:按置信度前top_n个，输出类别、置信度，
例如
[(18, 0.9916), (19, 0.0105), (15, 0.0038), (1, 0.0018), (5, 0.0016), (13, 0.0011)]
'''
bboxes = utils.postprocess_boxes(pred_bbox, (org_h, org_w), self.input_size, self.score_cls_threshold)
classes_in_img = list(set(bboxes[:, 5]))
best_bboxes = {}
for cls in classes_in_img:
cls_mask = (bboxes[:, 5] == cls)
cls_bboxes = bboxes[cls_mask]
best_score = 0
for i in range(len(cls_bboxes)):
if cls_bboxes[i][-2] > best_score:
best_score = cls_bboxes[i][-2]
if int(cls) not in best_bboxes.keys():
best_bboxes[int(cls)] = round(best_score, 4)
best_bboxes = sorted(best_bboxes.items(), key=lambda best_bboxes: best_bboxes[1], reverse=True)
return best_bboxes[:top_n]
def predict(self, image):
org_image = np.copy(image)
org_h, org_w, _ = org_image.shape
image_data = utils.image_preporcess(image,
[self.input_size, self.input_size])
image_data = image_data[np.newaxis, ...]
pred_sbbox, pred_mbbox, pred_lbbox, layer_n = self.sess.run(
[
self.pred_sbbox, self.pred_mbbox, self.pred_lbbox,
self.layer_num
],
feed_dict={
self.input: image_data,
self.trainable: False
})
pred_bbox = np.concatenate([
np.reshape(pred_sbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_mbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_lbbox, (-1, 5 + self.num_classes))
],
axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, (org_h, org_w),
self.input_size, self.score_threshold)
bboxes = utils.nms(bboxes, self.iou_threshold)
return org_h, org_w, pred_bbox, bboxes, layer_n
def predict(self, image):
'''
预测结果
:param image: 图片数据，shape为[800,600,3]
:return:
bboxes：食材检测预测框结果，格式为：[x_min, y_min, x_max, y_max, probability, cls_id],
layer_n[0]：烤层检测结果，0：最下层、1：中间层、2：最上层、3：其他
'''
org_image = np.copy(image)
org_h, org_w, _ = org_image.shape
image_data = utils.image_preporcess(image, [self.input_size, self.input_size])
image_data = image_data[np.newaxis, ...]
pred_sbbox, pred_mbbox, pred_lbbox, layer_n = self.sess.run(
[self.pred_sbbox, self.pred_mbbox, self.pred_lbbox, self.layer_num],
feed_dict={
self.input: image_data,
self.trainable: False
}
)
pred_bbox = np.concatenate([np.reshape(pred_sbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_mbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_lbbox, (-1, 5 + self.num_classes))], axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, (org_h, org_w), self.input_size, self.score_threshold)
bboxes = utils.nms(bboxes, self.iou_threshold)
return bboxes, layer_n[0]
def predict(self, image):
'''
:param image: 图片数据
:return: std_result：食材结果[name, flag, num_label]
, layer_n:烤层结果，0：下层、1：中层、2：上层、3：其他
'''
org_image = np.copy(image)
org_h, org_w, _ = org_image.shape
image_data = utils.image_preporcess(image,
[self.input_size, self.input_size])
image_data = image_data[np.newaxis, ...]
pred_sbbox, pred_mbbox, pred_lbbox, layer_ = self.sess.run(
[
self.pred_sbbox, self.pred_mbbox, self.pred_lbbox,
self.layer_num
],
feed_dict={
self.input: image_data,
self.trainable: False
})
pred_bbox = np.concatenate([
np.reshape(pred_sbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_mbbox, (-1, 5 + self.num_classes)),
np.reshape(pred_lbbox, (-1, 5 + self.num_classes))
],
axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, (org_h, org_w),
self.input_size, 0.45)
bboxes = utils.nms(bboxes, self.iou_threshold)
bboxes_pr = bboxes # 检测框结果
layer_n = layer_[0] # 烤层结果
num_label = len(bboxes_pr)
# 未检测食材
if num_label == 0:
return bboxes_pr, layer_n
# 检测到一个食材
elif num_label == 1:
return bboxes_pr, layer_n
# 检测到多个食材
else:
same_label = True
for i in range(num_label):
if i == (num_label - 1):
break
if bboxes_pr[i][5] == bboxes_pr[i + 1][5]:
continue
else:
same_label = False
# 多个食材，同一标签
if same_label:
sumProb = 0.
for i in range(num_label):
sumProb += bboxes_pr[i][4]
avrProb = sumProb / num_label
bboxes_pr[0][4] = avrProb
return bboxes_pr, layer_n
# 多个食材，非同一标签
else:
sumProb = 0.
problist = list(map(lambda x: x[4], bboxes_pr))
labellist = list(map(lambda x: x[5], bboxes_pr))
labeldict = {}
for key in labellist:
labeldict[key] = labeldict.get(key, 0) + 1
# 按同种食材label数量降序排列
s_labeldict = sorted(labeldict.items(),
key=lambda x: x[1],
reverse=True)
# n_name = len(s_labeldict)
name1 = s_labeldict[0][0]
num_name1 = s_labeldict[0][1]
# 数量最多label对应的食材占比0.7以上
if num_name1 / num_label > 0.7:
num_label0 = []
for i in range(num_label):
if name1 == bboxes_pr[i][5]:
num_label0.append(bboxes_pr[i])
for i in range(len(num_label0)):
sumProb += num_label0[i][4]
avrProb = sumProb / num_label
num_label0[0][4] = avrProb
return num_label0, layer_n
# 按各个label的probability降序排序
else:
bboxes_pr = sorted(bboxes_pr,
key=lambda x: x[4],
reverse=True)
return bboxes_pr, layer_n
pred_sbbox, pred_mbbox, pred_lbbox = sess.run(
[return_tensors[2], return_tensors[3], return_tensors[4]],
feed_dict={
return_tensors[0]: image_data,
return_tensors[1]: False
})
pred_bbox = np.concatenate([
np.reshape(pred_sbbox, (-1, 5 + num_classes)),
np.reshape(pred_mbbox, (-1, 5 + num_classes)),
np.reshape(pred_lbbox, (-1, 5 + num_classes))
],
axis=0)
bboxes = utils.postprocess_boxes(pred_bbox, frame_size, input_size,
0.3)
bboxes = utils.nms(bboxes, 0.45, method='nms')
image_ = utils.draw_bbox(frame, bboxes)
curr_time = time.time()
exec_time = curr_time - prev_time
result = np.asarray(image_)
info = "time: %.2f ms" % (1000 * exec_time)
cv2.namedWindow("result", cv2.WINDOW_AUTOSIZE)
result_ = cv2.cvtColor(image_, cv2.COLOR_RGB2BGR)
cv2.imshow("result", result_)
key = cv2.waitKey(1) & 0xFF
if key == ord(" "): # 空格键暂停
cv2.waitKey(0)
if key == ord("q"):
