Tensorflow-混合不同的图形进行多类检测

时间:2018-10-23 07:29:44

标签: python-3.x opencv tensorflow

我需要同时在tensorflow上使用2个图,以检测不同类别的对象。 所以基本上首先我加载第一个模型并执行检测:

# Name of the directory containing the object detection module we're using
#MODEL_NAME = 'faster_rcnn_inception_v2_coco_2018_01_28'
MODEL_NAME = 'modelX_model'
LABEL_DIRE = 'modelX_model'
CWD_PATH = os.getcwd()
PATH_TO_CKPT = os.path.join(CWD_PATH,MODEL_NAME,'frozen_inference_graph.pb')
PATH_TO_LABELS = os.path.join(CWD_PATH,LABEL_DIRE,'labelmap.pbtxt')
NUM_CLASSES = 1
label_map_X = label_map_util.load_labelmap(PATH_TO_LABELS)
categories_X = label_map_util.convert_label_map_to_categories(label_map_X, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index_X = label_map_util.create_category_index(categories_X)


# Load the Tensorflow model into memory.
detection_graph_X = g_1
with detection_graph_X.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

    sess = tf.Session(graph=detection_graph_X)

# Input tensor is the image
image_tensor_X = detection_graph_X.get_tensor_by_name('image_tensor:0')
detection_boxes_X = detection_graph_X.get_tensor_by_name('detection_boxes:0')
detection_scores_X = detection_graph_X.get_tensor_by_name('detection_scores:0')
detection_classes_X = detection_graph_X.get_tensor_by_name('detection_classes:0')

# Number of objects detected
num_detections_X = detection_graph_X.get_tensor_by_name('num_detections:0')

image = cv2.imread(PATH_TO_IMAGE)
resize = cv2.resize(image, (800, 600), interpolation = cv2.INTER_LINEAR) 
image_expanded = np.expand_dims(image, axis=0)

# Perform the actual detection by running the model with the image as input
(boxes, scores, classes, num) = sess.run(
    [detection_boxes_X, detection_scores_X, detection_classes_X, num_detections_X],
    feed_dict={image_tensor_X: image_expanded})

vis_util.visualize_boxes_and_labels_on_image_array(
    image,
    np.squeeze(boxes),
    np.squeeze(classes).astype(np.int32),
    np.squeeze(scores),
    category_index_X,
    use_normalized_coordinates=True,
    line_thickness=8,
    min_score_thresh=0.80)

然后我加载第二个模型并执行检测:

# Name of the directory containing the object detection module we're using
MODEL_NAME = 'faster_rcnn_inception_v2_coco_2018_01_28'
LABEL_DIRE = 'generic_model'
CWD_PATH = os.getcwd()
PATH_TO_CKPT = os.path.join(CWD_PATH,MODEL_NAME,'frozen_inference_graph.pb')
PATH_TO_LABELS = os.path.join(CWD_PATH,LABEL_DIRE,'labelmap.pbtxt')
NUM_CLASSES = 1
label_map_gen = label_map_util.load_labelmap(PATH_TO_LABELS)
categories_gen = label_map_util.convert_label_map_to_categories(label_map_gen, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index_gen = label_map_util.create_category_index(categories_gen)

# Load the Tensorflow model into memory.
detection_graph_gen = g_2
with detection_graph_gen.as_default():
    od_graph_def_2 = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def_2.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def_2, name='')

    sess = tf.Session(graph=detection_graph_gen)

# Input tensor is the image
image_tensor_gen = detection_graph_gen.get_tensor_by_name('image_tensor:0')
detection_boxes_gen = detection_graph_gen.get_tensor_by_name('detection_boxes:0')
detection_scores_gen = detection_graph_gen.get_tensor_by_name('detection_scores:0')
detection_classes_gen = detection_graph_gen.get_tensor_by_name('detection_classes:0')

# Number of objects detected
num_detections_gen = detection_graph_gen.get_tensor_by_name('num_detections:0')

#image = cv2.imread(PATH_TO_IMAGE)
#image_expanded = np.expand_dims(image, axis=0)

# Perform the actual detection by running the model with the image as input
(boxes, scores, classes, num) = sess.run(
    [detection_boxes_gen, detection_scores_gen, detection_classes_gen, num_detections_gen],
    feed_dict={image_tensor_gen: image_expanded})

vis_util.visualize_boxes_and_labels_on_image_array(
    image,
    np.squeeze(boxes),
    np.squeeze(classes).astype(np.int32),
    np.squeeze(scores),
    category_index_X,
    use_normalized_coordinates=True,
    line_thickness=28,
    min_score_thresh=0.80)

最后,我使用打开的简历显示结果。

# All the results have been drawn on image. Now display the image.
cv2.imshow('Object detector', image)

我的问题是我无法区分颜色框和在每个模型Labelmap.pbtxt中指定的类:基本上检测似乎可行,但我无法在Cv2框中添加正确的对象命名。

您有什么建议吗? 是否有更好的方法同时使用2个图形?

Thx

0 个答案:

没有答案