Created
July 31, 2017 08:16
-
-
Save datitran/90cb3accb8d35148ab3d4d4f9360b0d4 to your computer and use it in GitHub Desktop.
Object Recognition App with Video as Source
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import cv2 | |
import time | |
import numpy as np | |
import tensorflow as tf | |
from utils import FPS | |
from object_detection.utils import label_map_util | |
from object_detection.utils import visualization_utils as vis_util | |
CWD_PATH = os.getcwd() | |
# Path to frozen detection graph. This is the actual model that is used for the object detection. | |
MODEL_NAME = 'dat_model' | |
PATH_TO_CKPT = os.path.join(CWD_PATH, 'object_detection', MODEL_NAME, 'output_inference_graph.pb') | |
# List of the strings that is used to add correct label for each box. | |
PATH_TO_LABELS = os.path.join(CWD_PATH, 'object_detection', 'data', 'object-detection.pbtxt') | |
NUM_CLASSES = 1 | |
# Loading label map | |
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) | |
def detect_objects(image_np, sess, detection_graph): | |
# Expand dimensions since the model expects images to have shape: [1, None, None, 3] | |
image_np_expanded = np.expand_dims(image_np, axis=0) | |
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. | |
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. | |
scores = detection_graph.get_tensor_by_name('detection_scores:0') | |
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_detections) = sess.run( | |
[boxes, scores, classes, num_detections], | |
feed_dict={image_tensor: image_np_expanded}) | |
#print(boxes, scores, classes, num_detections) | |
# Visualization of the results of a detection. | |
vis_util.visualize_boxes_and_labels_on_image_array( | |
image_np, | |
np.squeeze(boxes), | |
np.squeeze(classes).astype(np.int32), | |
np.squeeze(scores), | |
category_index, | |
use_normalized_coordinates=True, | |
line_thickness=8, | |
min_score_thresh=0.5) | |
return image_np | |
if __name__ == '__main__': | |
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='') | |
sess = tf.Session(graph=detection_graph) | |
video_capture = cv2.VideoCapture('video.mp4') | |
fps = FPS().start() | |
frame_width = int(video_capture.get(3)) | |
frame_height = int(video_capture.get(4)) | |
# define video output | |
out = cv2.VideoWriter('outpy.avi', cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (frame_width, frame_height)) | |
count = 0 | |
while video_capture.isOpened(): | |
ret, frame = video_capture.read() | |
t = time.time() | |
detected_image = detect_objects(frame, sess, detection_graph) | |
fps.update() | |
count += 1 | |
if count % 100 == 0: | |
print(count) | |
# write to video file | |
out.write(detected_image) | |
cv2.imshow('Video', detected_image) | |
print('[INFO] elapsed time: {:.2f}'.format(time.time() - t)) | |
if cv2.waitKey(1) & 0xFF == ord('q'): | |
break | |
fps.stop() | |
print('[INFO] elapsed time (total): {:.2f}'.format(fps.elapsed())) | |
print('[INFO] approx. FPS: {:.2f}'.format(fps.fps())) | |
video_capture.release() | |
sess.close() | |
cv2.destroyAllWindows() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment