将检测坐标输入到对象跟踪器

时间:2019-07-23 09:35:50

标签: python opencv tensorflow object-detection object-detection-api

我正在进行多对象跟踪,正在使用TensorFlow API生成检测。我设法对其进行了一些修改,以使其返回检测到的对象的坐标,现在我想将坐标(边界框)馈送到对象跟踪器(CRST或KCF)。 然而,同时运行检测和跟踪将在计算上过于昂贵。 还有其他方法可以通过坐标或暂停检测吗? 下面是检测代码。 在此链接中,跟踪代码为https://github.com/spmallick/learnopencv/blob/master/MultiObjectTracker/multiTracker.py


import numpy as np
import os
import six.moves.urllib as urllib
import sys
sys.path.insert(0,r'C:\Users\Ahmed.DESKTOP-KJ6U1BJ\.spyder-py3\TensorFlow\models\research\object_detection')
import tarfile
import tensorflow as tf
import zipfile
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
import cv2
import imutils
from protos import string_int_label_map_pb2
from utils import visualization_utils2 as vis_util



def scale(bbox, width, height):
    x = int(bbox[0]*width)
    y = int(bbox[1]*height)
    w = int(bbox[2]*width)
    h = int(bbox[3]*height)
    return (x,y,w,h)

W = 800
H  = 600



videopath = "file:///C:/Users/Ahmed.DESKTOP-KJ6U1BJ/.spyder-py3/soccer4.mp4"
cap = cv2.VideoCapture(videopath)
# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")
# # Model preparation 
# Any model exported using the `export_inference_graph.py` tool can be loaded here simply by changing `PATH_TO_CKPT` to point to a new .pb file.  
# By default we use an "SSD with Mobilenet" model here. See the [detection model zoo](https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md) for a list of other models that can be run out-of-the-box with varying speeds and accuracies.

# What model to download.
MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'
# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = r'C:\Users\Ahmed.DESKTOP-KJ6U1BJ\.spyder-py3\TensorFlow\models\research\object_detection\data\mscoco_label_map.pbtxt'
NUM_CLASSES = 90

# ## Download Model ( uncomment if the model isn't downloaded / comment if you alredy have the model)
"""
opener = urllib.request.URLopener()
opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
tar_file = tarfile.open(MODEL_FILE)
for file in tar_file.getmembers():
  file_name = os.path.basename(file.name)
  if 'frozen_inference_graph.pb' in file_name:
    tar_file.extract(file, os.getcwd())
"""
# ## Load a (frozen) Tensorflow model into memory.

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='')

# ## Loading label map
# Label maps map indices to category names, so that when our convolution network predicts `5`, we know that this corresponds to `airplane`.  Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine

import label_map_util

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)

# # Detection


with detection_graph.as_default():
  with tf.Session(graph=detection_graph) as sess:
    while True :

      ret, image_np = cap.read()
      # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
      image_np_expanded = np.expand_dims(image_np, axis=0)
      # 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')
      # Actual detection.
      (boxes, scores, classes, num) = sess.run(
          [detection_boxes, detection_scores, detection_classes, num_detections],
          feed_dict={image_tensor: image_np_expanded})
      # Visualization of the results of a detection.
      boxes2 = np.squeeze(boxes)
      max_boxes_to_draw =boxes2.shape[0]
      scores2 = np.squeeze(scores)
      min_score_thresh=0.7
      classes2 = np.squeeze(classes).astype(np.int32)
      for i in range(min(max_boxes_to_draw, boxes2.shape[0])):
        if boxes2 is None or scores2[i] > min_score_thresh:
          class_name = category_index[classes2[i]]['name']
          print ("This box is gonna get used", scale(boxes2[i], W ,  H),  class_name)

      cv2.imshow('Object Detection',cv2.resize(image_np,(800,600)))
      k = cv2.waitKey(1) & 0xff
      if k == 27:
        cv2.destroyAllWindows()
        cap.release()    

cv2.destroyAllWindows()
cap.release

1 个答案:

答案 0 :(得分:0)

您可以在while True循环中使用简单的计数器对帧进行计数,并在if之前使用session.run语句“暂停”检测,例如:

frame_count = 0
with detection_graph.as_default():
  with tf.Session(graph=detection_graph) as sess:
    while True :           
      ret, image_np = cap.read()
      #the first frame and every 10 frames do the detection
      if frame_count == 0:
         ###detection here
         #restart counter (from -10 to 0)
         frame_count = -10
      ##do tracking here
      frame_count += 1

通过这种方式,对第一帧然后每10帧进行一次实际检测,因此在其他9帧中,您可以执行所需的任何操作。