如何在tensorflow对象检测中仅获取特定类别的框

时间:2017-10-01 18:10:26

标签: tensorflow object-detection

object detection notebook演示了预训练和冻结张量流模型如何用于检测测试图像中的对象。

在这个notetook中,功能

from utils import visualization_utils as vis_util
vis_util.visualize_boxes_and_labels_on_image_array(
      image_np,
      np.squeeze(boxes),
      np.squeeze(classes).astype(np.int32),
      np.squeeze(scores),
      category_index,
      use_normalized_coordinates=True,
      line_thickness=8)

输出测试图像,其中在检测到的对象周围绘制框。

如何使用此功能仅绘制特定类别的框而不是category_index集中所有类别的框?也就是说,我怎样才能使这个函数只能在模型确定它们的对象周围绘制框。汽车?

1 个答案:

答案 0 :(得分:3)

这个答案对你来说可能为时已晚。但希望这会对其他人有所帮助。

object_detection 文件夹中,文件夹名称 utils 里面有python文件名 visualization_utils.py ,您必须编辑名为的函数此文件上的visualize_boxes_and_labels_on_image_array 。在此函数中,它在此添加之前调用名为draw_bounding_box_on_image_array的函数 display_str_list = box_to_display_str_map [box] 并添加if条件,查找对象类的内容。在此if条件内调用draw_bounding_box_on_image_array。(示例代码如下所示,您可以更改对象检测标签名称)

display_str_list = box_to_display_str_map[box]
if (("bottle" in display_str_list[0]) or ("person") in display_str_list[0]):
  draw_bounding_box_on_image_array(
    image,
    ymin,
    xmin,
    ymax,
    xmax,
    color= color,
    thickness= line_thickness,
    display_str_list=box_to_display_str_map[box],
    use_normalized_coordinates=use_normalized_coordinates)
  if keypoints is not None:
    draw_keypoints_on_image_array(
      image,
      box_to_keypoints_map[box],
      color= color,
      radius= line_thickness /.2,
      use_normalized_coordinates= use_normalized_coordinates)