目标检测——SSD预测及后处理
1. 加载model及训练权重
self.num_classes = len(self.class_names) + 1 self.model = SSD300((300,300,3), self.num_classes) self.model.load_weights(\'ssd_model.h5\', by_name=True, skip_mismatch=True) # self.model.summary() print(\'model loaded\')
2. 图片预处理
将图片缩放至(300, 300)大小,并进行与训练过程同样的预处理(不包括图像增强)输入网络,得到预测结果
photo = preprocess_input(np.reshape(photo,[1,self.model_image_size[0],self.model_image_size[1],3]))
preds = self.get_pred(photo).numpy()
3. 预测结果后处理
3.1 得到所有预测框的位置
def decode_boxes(self, mbox_loc, mbox_priorbox, variances): # 获得先验框的宽与高 prior_width = mbox_priorbox[:, 2] - mbox_priorbox[:, 0] prior_height = mbox_priorbox[:, 3] - mbox_priorbox[:, 1] # 获得先验框的中心点 prior_center_x = 0.5 * (mbox_priorbox[:, 2] + mbox_priorbox[:, 0]) prior_center_y = 0.5 * (mbox_priorbox[:, 3] + mbox_priorbox[:, 1]) # 真实框距离先验框中心的xy轴偏移情况 decode_bbox_center_x = mbox_loc[:, 0] * prior_width * variances[:, 0] decode_bbox_center_x += prior_center_x decode_bbox_center_y = mbox_loc[:, 1] * prior_height * variances[:, 1] decode_bbox_center_y += prior_center_y # 真实框的宽与高的求取 decode_bbox_width = np.exp(mbox_loc[:, 2] * variances[:, 2]) decode_bbox_width *= prior_width decode_bbox_height = np.exp(mbox_loc[:, 3] * variances[:, 3]) decode_bbox_height *= prior_height # 获取真实框的左上角与右下角 decode_bbox_xmin = decode_bbox_center_x - 0.5 * decode_bbox_width decode_bbox_ymin = decode_bbox_center_y - 0.5 * decode_bbox_height decode_bbox_xmax = decode_bbox_center_x + 0.5 * decode_bbox_width decode_bbox_ymax = decode_bbox_center_y + 0.5 * decode_bbox_height # 真实框的左上角与右下角进行堆叠 decode_bbox = np.concatenate((decode_bbox_xmin[:, None], decode_bbox_ymin[:, None], decode_bbox_xmax[:, None], decode_bbox_ymax[:, None]), axis=-1) # 防止超出0与1 decode_bbox = np.minimum(np.maximum(decode_bbox, 0.0), 1.0) return decode_bbox
decode_bbox 的维度为(8732, 4)
3.2 针对每一个类别将置信度高于一定阈值 (一般设为0.5) 的作为备选框
for c in range(self.num_classes): if c == background_label_id: continue c_confs = mbox_conf[i, :, c] c_confs_m = c_confs > confidence_threshold if len(c_confs[c_confs_m]) > 0: # 取出得分高于confidence_threshold的框 boxes_to_process = decode_bbox[c_confs_m] confs_to_process = c_confs[c_confs_m]
3.3 采取非极大值抑制过滤重合度较高的框
# 进行iou的非极大抑制 idx = tf.image.non_max_suppression(tf.cast(boxes_to_process,tf.float32), tf.cast(confs_to_process,tf.float32),self._top_k, iou_threshold=self._nms_thresh).numpy() # 取出在非极大抑制中效果较好的内容 good_boxes = boxes_to_process[idx] confs = confs_to_process[idx][:, None] # 将label、置信度、框的位置进行堆叠。 labels = c * np.ones((len(idx), 1)) c_pred = np.concatenate((labels, confs, good_boxes), axis=1) # 添加进result里 results[-1].extend(c_pred)
c_pred 的维度为(n , 6), n为c类别框的个数
3.4 按置信度排序,选取top_k个框
results[-1] = np.array(results[-1]) argsort = np.argsort(results[-1][:, 1])[::-1] results[-1] = results[-1][argsort] # 选出置信度最大的keep_top_k个 results[-1] = results[-1][:keep_top_k] return results
4. 将预测结果(框的位置)映射回原图像大小
def ssd_correct_boxes(top, left, bottom, right, input_shape, image_shape): new_shape = image_shape*np.min(input_shape/image_shape) offset = (input_shape-new_shape)/2./input_shape scale = input_shape/new_shape box_yx = np.concatenate(((top+bottom)/2,(left+right)/2),axis=-1) box_hw = np.concatenate((bottom-top,right-left),axis=-1) box_yx = (box_yx - offset) * scale box_hw *= scale box_mins = box_yx - (box_hw / 2.) box_maxes = box_yx + (box_hw / 2.) boxes = np.concatenate([ box_mins[:, 0:1], box_mins[:, 1:2], box_maxes[:, 0:1], box_maxes[:, 1:2] ],axis=-1) print(np.shape(boxes)) boxes *= np.concatenate([image_shape, image_shape],axis=-1) return boxes
5.在原图像上画框
for i, c in enumerate(top_label_indices): predicted_class = self.class_names[int(c)-1] score = top_conf[i] top, left, bottom, right = boxes[i] top = top - 5 left = left - 5 bottom = bottom + 5 right = right + 5 top = max(0, np.floor(top + 0.5).astype(\'int32\')) left = max(0, np.floor(left + 0.5).astype(\'int32\')) bottom = min(np.shape(image)[0], np.floor(bottom + 0.5).astype(\'int32\')) right = min(np.shape(image)[1], np.floor(right + 0.5).astype(\'int32\')) # 画框框 label = \'{} {:.2f}\'.format(predicted_class, score) draw = ImageDraw.Draw(image) label_size = draw.textsize(label, font) label = label.encode(\'utf-8\') print(label) if top - label_size[1] >= 0: text_origin = np.array([left, top - label_size[1]]) else: text_origin = np.array([left, top + 1]) for i in range(thickness): draw.rectangle( [left + i, top + i, right - i, bottom - i], outline=self.colors[int(c)-1]) draw.rectangle( [tuple(text_origin), tuple(text_origin + label_size)], fill=self.colors[int(c)-1]) draw.text(text_origin, str(label,\'UTF-8\'), fill=(0, 0, 0), font=font) del draw return image