Python label_map_util.create_category_index_from_labelmap方法代码示例

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # Initialise TF graph and category index
        self._graph = tf.Graph()
        with self._graph.as_default():
            od_graph_def = tf.compat.v1.GraphDef()
            with tf.compat.v1.gfile.GFile(str(self.model_path), 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')
            self._sess = tf.compat.v1.Session(graph=self._graph, config=self.session_config)
        self.category_index = lm_util.create_category_index_from_labelmap(
            self.labels_path,
            use_display_name=True)

        # Variables used in async mode
        if self.run_async:
            self._buffer = None
            # Initialise frame capture thread
            self._capture_thread = threading.Thread(target=self._capture)
            self._capture_thread.daemon = True
            self._thread_lock = threading.Lock()
            self._capture_thread.start() 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def _load_graph(self):
        self.detection_graph = tf.Graph()
        with self.detection_graph.as_default() as default_graph:
            od_graph_def = tf.GraphDef()
            with tf.gfile.GFile(self.frozen_graph_path, 'rb') as fid:
                serialized_graph = fid.read()
                od_graph_def.ParseFromString(serialized_graph)
                tf.import_graph_def(od_graph_def, name='')

        config = tf.ConfigProto()
        config.gpu_options.allow_growth = True
        config.log_device_placement = True

        self.category_index = label_map_util.create_category_index_from_labelmap(self.label_path, use_display_name=True)
        self.session = tf.Session(config=config, graph=default_graph)
        self.global_graph = default_graph 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def test_create_category_index_from_labelmap(self):
    label_map_string = """
      item {
        id:2
        name:'cat'
      }
      item {
        id:1
        name:'dog'
      }
    """
    label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
    with tf.gfile.Open(label_map_path, 'wb') as f:
      f.write(label_map_string)

    category_index = label_map_util.create_category_index_from_labelmap(
        label_map_path)
    self.assertDictEqual({
        1: {
            'name': u'dog',
            'id': 1
        },
        2: {
            'name': u'cat',
            'id': 2
        }
    }, category_index) 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def test_create_category_index_from_labelmap_display(self):
    label_map_string = """
      item {
        id:2
        name:'cat'
        display_name:'meow'
      }
      item {
        id:1
        name:'dog'
        display_name:'woof'
      }
    """
    label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
    with tf.gfile.Open(label_map_path, 'wb') as f:
      f.write(label_map_string)

    self.assertDictEqual({
        1: {
            'name': u'dog',
            'id': 1
        },
        2: {
            'name': u'cat',
            'id': 2
        }
    }, label_map_util.create_category_index_from_labelmap(
        label_map_path, False))

    self.assertDictEqual({
        1: {
            'name': u'woof',
            'id': 1
        },
        2: {
            'name': u'meow',
            'id': 2
        }
    }, label_map_util.create_category_index_from_labelmap(label_map_path)) 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def main(conf, conf_path, label_path, **kwargs):
    """Main function for receiver
    
    Args:
        conf (dict): Configuration file
        conf_path (str): Configuration path (plugins)
    
    Yields:
        bool: Detection successful
    """
    # List of the strings that is used to add correct label for each box.
    PATH_TO_LABELS = os.path.abspath(label_path)
    category_index = label_map_util.create_category_index_from_labelmap(
        PATH_TO_LABELS, use_display_name=True)

    detection_model = load_model(conf['Tensorflow']['ModelUrl'])

    # Client Plugins
    loaded_plugins = load_plugins(plugins=conf['Plugins']['Enabled'].split(
        ','), conf_path=conf_path+'/plugins.d')

    # Start loop
    for res in receive(category_index,
                       detection_model,
                       conf['ZmqCamera']['IP'],
                       conf['ZmqCamera']['Port'],
                       conf['ZmqCamera']['Protocol'],
                       int(conf['ZmqCamera']['Pattern']),
                       float(conf['Detection']['min_detections']),
                       float(conf['Detection']['min_confidence']),
                       server_plugins=loaded_plugins,
                       **kwargs):
        logger.debug('Received signal')
        if kwargs.get('use_sender_thread', False):
            send_async_messages(loaded_plugins)
        else:
            send_messages(loaded_plugins)
        # For downstream
        yield res 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def test_create_category_index_from_labelmap(self):
    label_map_string = """
      item {
        id:2
        name:'cat'
      }
      item {
        id:1
        name:'dog'
      }
    """
    label_map_path = os.path.join(self.get_temp_dir(), 'label_map.pbtxt')
    with tf.gfile.Open(label_map_path, 'wb') as f:
      f.write(label_map_string)

    category_index = label_map_util.create_category_index_from_labelmap(
        label_map_path)
    self.assertDictEqual({
        1: {
            'name': 'dog',
            'id': 1
        },
        2: {
            'name': 'cat',
            'id': 2
        }
    }, category_index) 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def detect_image(image_path):
    # load label map
    category_index = label_map_util.create_category_index_from_labelmap(
        PATH_TO_LABELS)

    # load detection graph
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')

    # define input/output tensors
    image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
    detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
    detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
    detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
    num_detections = detection_graph.get_tensor_by_name('num_detections:0')

    # load input image
    img = cv2.imread(image_path)
    if img is None:
        sys.exit('failed to load image: %s' % image_path)
    img = img[..., ::-1]  # BGR to RGB

    # run inference
    with detection_graph.as_default():
        with tf.Session() as sess:
            boxes, scores, classes, _ = sess.run(
                [detection_boxes, detection_scores, detection_classes, num_detections],
                feed_dict={image_tensor: np.expand_dims(img, 0)})

    # draw the results of the detection
    vis_util.visualize_boxes_and_labels_on_image_array(
        img,
        np.squeeze(boxes),
        np.squeeze(classes).astype(np.int32),
        np.squeeze(scores),
        category_index,
        use_normalized_coordinates=True,
        line_thickness=6,
        min_score_thresh=0.3)

    # save the output image
    img = img[..., ::-1]  # RGB to BGR
    cv2.imwrite(OUTPUT_PATH, img)

    print('Output has been written to %s\n' % OUTPUT_PATH) 

# 需要导入模块: from object_detection.utils import label_map_util [as 别名]
# 或者: from object_detection.utils.label_map_util import create_category_index_from_labelmap [as 别名]
def vis_detection_result(graph,image_path,output_image_path):
    with graph.as_default():
        ops=tf.get_default_graph().get_operations()
        all_tensor_names={output.name for op in ops for output in op.outputs}
        tensor_dict={}
        for key in [
            'num_detections','detection_boxes','detection_scores',
            'detection_classes','detection_masks'
        ]:
            tensor_name=key+':0'
            if tensor_name in all_tensor_names:
                tensor_dict[key]=tf.get_default_graph().get_tensor_by_name(tensor_name)

        image_tensor=tf.get_default_graph().get_tensor_by_name('image_tensor:0')

        with tf.Session() as sess:
            print('get in the session')
            image = util.data_preprocessing(image_path,target_size=640)
            image_np = np.expand_dims(image, axis=0)
            output_dict=sess.run(tensor_dict,feed_dict={image_tensor:image_np})
            # print(output_dict)
            # all outputs are float32 numpy arrays, so convert types as appropriate
            output_dict['num_detections'] = int(output_dict['num_detections'][0])
            output_dict['detection_classes'] = output_dict[
                'detection_classes'][0].astype(np.int64)
            output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
            output_dict['detection_scores'] = output_dict['detection_scores'][0]
            #print(output_dict)
            # return output_dict
            print('output_dict[\'detection_boxes\'] shape is {}'.format(output_dict['detection_boxes'].shape))
            print('output_dict[\'detection_scores\'] shape is {}'.format(output_dict['detection_scores'].shape))

            category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)

            image=vis_util.visualize_boxes_and_labels_on_image_array(
                image,
                output_dict['detection_boxes'],
                output_dict['detection_classes'],
                output_dict['detection_scores'],
                category_index,
                instance_masks=output_dict.get('detection_masks'),
                use_normalized_coordinates=True,
                line_thickness=3,min_score_thresh=0.3)

            plt.imsave(output_image_path,image)

            sess.close()