Необходимо знать время обнаружения объекта на основе классов (меток) - PullRequest
Новая
       19

Необходимо знать время обнаружения объекта на основе классов (меток)

0 голосов
/ 10 апреля 2019

Я работаю над API-интерфейсом обнаружения объектов Tensorflow для обнаружения определенных объектов из видеофайла. Я разработал код для обнаружения объекта в видео с помощью Opencv. Теперь я хочу узнать время обнаружения объекта на основе классов. я использовал 3 разных класса на этикетке. 

 import cv2
 import time
 import numpy as np
 import os
 import six.moves.urllib as urllib
 import sys
 import tensorflow as tf
 import matplotlib.pyplot as plt
 from matplotlib import pyplot as plt

 from utils import visualization_utils as vis_util
 from utils import label_map_util
 from multiprocessing.dummy import Pool as ThreadPool      



PATH_TO_CKPT = 'output_inference_graph/frozen_inference_graph.pb'


PATH_TO_LABELS = os.path.join('ssd_mobilenet_v1', 'label_map.pbtxt')

NUM_CLASSES = 3

sys.path.append("..")

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = 
label_map_util.convert_label_map_to_categories(label_map, 
max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
 IMAGE_SIZE = (12, 8)
   count = 0


def detect_in_video():

# VideoWriter is the responsible of creating a copy of the video
# used for the detections but with the detections overlays. Keep in
# mind the frame size has to be the same as original video.
out = cv2.VideoWriter('out_videos/l.avi', cv2.VideoWriter_fourcc(
    'M', 'J', 'P', 'G'), 20, (1280, 720))

detection_graph = tf.Graph()
with detection_graph.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='')

label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(
    label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)

with detection_graph.as_default():
    with tf.Session(graph=detection_graph) as sess:
        # Definite input and output Tensors for detection_graph
        image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
        # Each box represents a part of the image where a particular object
        # was detected.
        detection_boxes = detection_graph.get_tensor_by_name(
            'detection_boxes:0')
        # Each score represent how level of confidence for each of the objects.
        # Score is shown on the result image, together with the class
        # label.
        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')
        cap = cv2.VideoCapture('test_videos/end.mp4')

        while(cap.isOpened()):
            # Read the frame
            ret, frame = cap.read()

            # Recolor the frame. By default, OpenCV uses BGR color space.
            # This short blog post explains this better:
            # https://www.learnopencv.com/why-does-opencv-use-bgr-color-format/
            #color_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)





            image_np_expanded = np.expand_dims(frame, axis=0)

            # Actual detection.
            start_time = time.time()
            (boxes, scores, classes, num) = sess.run(
                [detection_boxes, detection_scores,
                    detection_classes, num_detections],
                feed_dict={image_tensor: image_np_expanded})
            print('Iteration : %.3f sec'%(time.time()-start_time))
             # Visualization of the results of a detection.
            # note: perform the detections using a higher threshold
            #if frame != color_frame:
             #   cv2.imwrite("te.jpg",color_frame)

            vis_util.visualize_boxes_and_labels_on_image_array(
                frame,
                np.squeeze(boxes),
                np.squeeze(classes).astype(np.int32),
                np.squeeze(scores),
                category_index,
                use_normalized_coordinates=True,
                line_thickness=8,
                min_score_thresh=.85)

            output_rgb = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)


            out.write(output_rgb)

            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

        out.release()
        cap.release()
        cv2.destroyAllWindows()


def main():
detect_in_video()


if __name__ == '__main__':
main()

Я ожидал, что распечатает и сохранит время обнаружения с именем метки

...