How to get only boxes for specific categories in tensorflow object detection

Question

The object detection notebook demonstrates, how pretrained and frozen tensorflow models can be used to detect objects in test images.

In this notetook, the function

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)

outputs test images, where boxes are drawn around detected objects.

How can I use this function to draw only boxes of a specific category instead of boxes of all categories in the category_index set? I.e., how can I get this function to only draw boxes around objects of which the model is sure they are e.g. cars?

Answer 1

This answer may be too late for you. but hope this will help to others.

Inside your object_detection folder has folder name utils inside that has python file name visualization_utils.py you have to edit function named visualize_boxes_and_labels_on_image_array on this file. In this function it call function named draw_bounding_box_on_image_array before this add display_str_list = box_to_display_str_map[box] and add if condition for what will you looking object class. inside this if condition call the draw_bounding_box_on_image_array.(example code given bellow and you can change object detection label names)

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)