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

Question

该物体检测笔记本演示，如何预先训练和冷冻tensorflow模型可以用来检测测试图像的对象。

在此通知中，功能

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集合中所有类别的框？即，如何获得此功能以仅在模型确定其为例如汽车的对象周围绘制框？

Answer 1

对于您来说，这个答案可能为时已晚。但希望这对其他人有帮助。

在object_detection文件夹中，文件夹名为utils，在其中具有python文件名visualization_utils.py，您必须在此文件上编辑名为visualize_boxes_and_labels_on_image_array的函数。在此函数中，它在添加display_str_list = box_to_display_str_map [box]之前调用名为draw_bounding_box_on_image_array的函数， 并为您要查找的对象类添加条件。在此条件条件内调用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)