Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- import multiprocessing as mp
- import time, logging, itertools
- import cv2 as cv
- import numpy as np
- import os
- import six.moves.urllib as urllib
- import sys
- import tarfile
- import tensorflow as tf
- import zipfile
- from distutils.version import StrictVersion
- from collections import defaultdict
- from io import StringIO
- from matplotlib import pyplot as plt
- from PIL import Image
- from object_detection.utils import ops as utils_ops
- from object_detection.utils import label_map_util
- from object_detection.utils import visualization_utils as vis_util
- def run_inference_for_single_image(sess, image, graph):
- with graph.as_default():
- # Get handles to input and output tensors
- 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)
- if 'detection_masks' in tensor_dict:
- # The following processing is only for single image
- detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
- detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
- # Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
- real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
- detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
- detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
- detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
- detection_masks, detection_boxes, image.shape[1], image.shape[2])
- detection_masks_reframed = tf.cast(
- tf.greater(detection_masks_reframed, 0.5), tf.uint8)
- # Follow the convention by adding back the batch dimension
- tensor_dict['detection_masks'] = tf.expand_dims(
- detection_masks_reframed, 0)
- image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
- # Run inference
- output_dict = sess.run(tensor_dict,
- feed_dict={image_tensor: image})
- # 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.uint8)
- output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
- output_dict['detection_scores'] = output_dict['detection_scores'][0]
- if 'detection_masks' in output_dict:
- output_dict['detection_masks'] = output_dict['detection_masks'][0]
- return output_dict
- # Holds a frame (numpy array) and creates current timestamp and autoincremental sequence number
- class frame_container:
- counter = itertools.count()
- def __init__(self, frame):
- self.frame = frame
- self.timestamp = time.time()
- self.seq = next(self.counter)
- # Plays the last frame generated by the object detection model
- def video_player(q):
- while True:
- if not q.empty():
- frame = q.get()
- if debug == 1: print("[VP] Received frame #{}".format(frame.seq))
- else:
- continue
- cv.imshow('SSD Output', frame.frame)
- if cv.waitKey(1000 // fps) & 0xFF == ord('q'):
- break
- # Puts frames into the model process queue, downsampling FPS from 30 to 5
- def producer(q):
- cap = cv.VideoCapture('road.mp4')
- n_frame = 0
- while (cap.isOpened()):
- ret, frame = cap.read()
- if not ret:
- cap = cv.VideoCapture('road.mp4')
- continue
- dims = frame.shape[:2]
- max_side = np.max(dims)
- res = cv.resize(frame, None, fx=500 / max_side, fy=500 / max_side, interpolation=cv.INTER_CUBIC)
- if not q.full():
- if n_frame == 0:
- frame = frame_container(res)
- q.put(frame, False)
- # cv.imshow('Original Video', frame.frame)
- if cv.waitKey(1000 // fps) & 0xFF == ord('q'):
- break
- if debug == 1: print("[P] Produced frame #{}".format(frame.seq))
- n_frame = (n_frame + 1) % 6
- cap.release()
- cv.destroyAllWindows()
- # Waits for #batch_size frames and runs them through the object detection model, sending output to the video player process
- def consumer(q, output_q):
- # What model to download.
- MODEL_NAME = 'ssd_mobilenet_v1_coco_2018_01_28'
- 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_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'
- # List of the strings that is used to add correct label for each box.
- PATH_TO_LABELS = os.path.abspath('../data/mscoco_label_map.pbtxt')
- # Uncomment to download new 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())
- 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='')
- category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
- model_output_imgs = []
- input_images = []
- inference_samples_num = 4
- beta = 1 - 1 / inference_samples_num
- inference_avg = -1
- inference_current = -1
- drift = 0
- with tf.Session(graph=detection_graph) as sess:
- while True:
- if not q.empty():
- if debug == 1: print("[C] Drift: {}".format(drift))
- frame = q.get()
- input_images.append(frame)
- # Drift is how much behind schedule the current frame is.
- # If more than one entire frame behind, skip this frame.
- if drift >= 1 / fps:
- drift -= 1 / fps
- input_images.pop(0)
- if debug >= 1: print("[C] Skipping frame #{}".format(frame.seq))
- else:
- drift += max(inference_current / batch_size - 1 / fps, 0)
- if len(input_images) >= batch_size:
- tick = time.time()
- if debug == 1: print("[C] Q, I, O, P size: ({}, {}, {}, {})".format(q.qsize(), len(input_images), len(model_output_imgs), output_q.qsize()))
- for i in range(batch_size):
- minitick = time.time()
- current_frame = input_images.pop(0)
- image_np = current_frame.frame
- # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
- image_np_expanded = np.expand_dims(image_np, axis=0)
- # Actual detection.
- output_dict = run_inference_for_single_image(sess, image_np_expanded, detection_graph)
- # Visualization of the results of a detection.
- vis_util.visualize_boxes_and_labels_on_image_array(
- image_np,
- 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=2)
- current_frame.frame = image_np
- minitock = time.time()
- output_q.put(current_frame, False)
- # Fake delay so that batch inference time is >2
- # if minitock - minitick < 2 / batch_size:
- # time.sleep(2 / batch_size - (time.time() - minitick))
- tock = time.time()
- inference_current = (tock - tick)
- if inference_avg < 0:
- inference_avg = inference_current
- inference_avg = beta * inference_avg + (1 - beta) * inference_current
- print("[C] Inference current batch: {}".format(inference_current))
- # print("[C] Inference average: {}".format(inference_avg))
- cv.destroyAllWindows()
- # logger = mp.log_to_stderr(logging.DEBUG)
- fps = 5
- batch_size = 8
- debug = 2
- # 0: inference time
- # 1: everything
- # 2: skipped frames
- consumer_in_q = mp.Queue(20)
- consumer_out_q = mp.Queue(20)
- #Producer process
- p_proc = mp.Process(target=producer, args=(consumer_in_q,))
- #Consumer process
- c_proc = mp.Process(target=consumer, args=(consumer_in_q, consumer_out_q))
- #Video player process
- player = mp.Process(target=video_player, args=(consumer_out_q,))
- p_proc.start()
- c_proc.start()
- player.start()
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement