Выбор классов обнаружения объектов Tensorflow SSD

Question

Я работаю над примером обнаружения тензорного объекта здесь он обнаруживает все 90 классов из набора данных Coco. Но я хочу обнаружить только два класса из этого. как это сделать?

Answer 1

Допустим, вы хотите обнаружить только мотоцикл и человека . внутри файла visualization_util.py, перейдите к def draw_bounding_box_on_image_array есть функция:

  draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,
                         thickness, display_str_list,
                         use_normalized_coordinates)

поместите этот вызов функции внутри условия, подобного этому

if (display_str_list[0][0:3]=="per" or display_str_list[0][0:3]=="mot"):

  draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,
                         thickness, display_str_list,
                         use_normalized_coordinates)

здесь «per» - первые три буквы «person», а «mot» - первые три буквы мотоцикла. Таким образом, вы можете обнаружить желаемый объект из всех других