diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2e3fb06 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +.idea/ +.idea/* diff --git a/__pycache__/detector.cpython-36.pyc b/__pycache__/detector.cpython-36.pyc new file mode 100644 index 0000000..5afbcac Binary files /dev/null and b/__pycache__/detector.cpython-36.pyc differ diff --git a/__pycache__/helpers.cpython-36.pyc b/__pycache__/helpers.cpython-36.pyc new file mode 100644 index 0000000..9b5ee22 Binary files /dev/null and b/__pycache__/helpers.cpython-36.pyc differ diff --git a/__pycache__/tracker.cpython-36.pyc b/__pycache__/tracker.cpython-36.pyc new file mode 100644 index 0000000..661cd9d Binary files /dev/null and b/__pycache__/tracker.cpython-36.pyc differ diff --git a/detector.py b/detector.py index eae382a..0077cb9 100644 --- a/detector.py +++ b/detector.py @@ -33,18 +33,18 @@ def __init__(self): self.detection_graph = tf.Graph() # configuration for possible GPU use - config = tf.ConfigProto() + config = tf.compat.v1.ConfigProto() config.gpu_options.allow_growth = True # load frozen tensorflow detection model and initialize # the tensorflow graph with self.detection_graph.as_default(): - od_graph_def = tf.GraphDef() - with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: + od_graph_def = tf.compat.v1.GraphDef() + with tf.compat.v1.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='') - self.sess = tf.Session(graph=self.detection_graph, config=config) + self.sess = tf.compat.v1.Session(graph=self.detection_graph, config=config) self.image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0') # Each box represents a part of the image where a particular object was detected. self.boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0') diff --git a/main.py b/main.py index 265c034..4b7c9dc 100644 --- a/main.py +++ b/main.py @@ -2,17 +2,18 @@ # -*- coding: utf-8 -*- """@author: kyleguan """ +import time import numpy as np import matplotlib.pyplot as plt import glob from moviepy.editor import VideoFileClip from collections import deque -from sklearn.utils.linear_assignment_ import linear_assignment +from scipy.optimize import linear_sum_assignment import helpers import detector -import tracker +from tracker import Tracker # Global variables to be used by funcitons of VideoFileClop frame_count = 0 # frame counter @@ -39,13 +40,15 @@ def assign_detections_to_trackers(trackers, detections, iou_thrd = 0.3): #trk = convert_to_cv2bbox(trk) for d,det in enumerate(detections): # det = convert_to_cv2bbox(det) - IOU_mat[t,d] = box_iou2(trk,det) + IOU_mat[t,d] = helpers.box_iou2(trk, det) # Produces matches # Solve the maximizing the sum of IOU assignment problem using the # Hungarian algorithm (also known as Munkres algorithm) - matched_idx = linear_assignment(-IOU_mat) + indices = linear_sum_assignment(-IOU_mat) + indices = np.asarray(indices) + matched_idx = np.transpose(indices) unmatched_trackers, unmatched_detections = [], [] for t,trk in enumerate(trackers): @@ -108,8 +111,7 @@ def pipeline(img): x_box.append(trk.box) - matched, unmatched_dets, unmatched_trks \ - = assign_detections_to_trackers(x_box, z_box, iou_thrd = 0.3) + matched, unmatched_dets, unmatched_trks = assign_detections_to_trackers(x_box, z_box, iou_thrd = 0.3) if debug: print('Detection: ', z_box) print('x_box: ', x_box) @@ -191,23 +193,11 @@ def pipeline(img): if __name__ == "__main__": det = detector.CarDetector() - - if debug: # test on a sequence of images - images = [plt.imread(file) for file in glob.glob('./test_images/*.jpg')] - - for i in range(len(images))[0:7]: - image = images[i] - image_box = pipeline(image) - plt.imshow(image_box) - plt.show() - - else: # test on a video file. - - start=time.time() - output = 'test_v7.mp4' - clip1 = VideoFileClip("project_video.mp4")#.subclip(4,49) # The first 8 seconds doesn't have any cars... - clip = clip1.fl_image(pipeline) - clip.write_videofile(output, audio=False) - end = time.time() - - print(round(end-start, 2), 'Seconds to finish') + start=time.time() + output = 'test_v7.mp4' + clip1 = VideoFileClip("project_video.mp4")#.subclip(4,49) # The first 8 seconds doesn't have any cars... + clip = clip1.fl_image(pipeline) + clip.write_videofile(output, audio=False) + end = time.time() + + print(round(end-start, 2), 'Seconds to finish') diff --git a/test_v7.mp4 b/test_v7.mp4 index 8680885..72b18fe 100644 Binary files a/test_v7.mp4 and b/test_v7.mp4 differ