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
集中所有类别的框?也就是说,我怎样才能使这个函数只能在模型确定它们的对象周围绘制框。汽车?
答案 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)