Skip to content
Permalink
 
 
Cannot retrieve contributors at this time
import numpy as np
import cv2
from skimage.feature import hog
from skimage.transform import resize
from scipy.ndimage.measurements import label
from vehicletracker.features import FeatureExtractor
from collections import deque
class VehicleTracker(object):
"""
Tracks surrounding vehicles in a series of consecutive frames.
"""
def __init__(self, scaler, classifier, first_frame):
"""
Initialises an instance.
Parameters
----------
scaler : SciPy scaler to apply to X.
classifier : Trained SciPy classifier for detecting vehicles.
first_frame : First video frame.
"""
self.scaler = scaler
self.classifier = classifier
self.frame_shape = first_frame.shape
self.detections_history = deque(maxlen=20)
def process(self, frame, draw_detections=True):
"""
Perform single frame processing and saves detected vehicles data.
Parameters
----------
frame : Current video frame.
draw_detections : Flag indicating if we need to highlight vehicles in the frame.
Returns
-------
Video frame
"""
self.detect_vehicles(frame)
if draw_detections:
for c in self.detections():
cv2.rectangle(frame, (c[0], c[1]), (c[2], c[3]), (0, 0, 255), 2)
return frame
def detections(self):
"""
Approximates surrounding vehicles based on a heatmap of last N detections.
Returns
-------
Boundaries of detected vehicles.
"""
detections, _ = self.merge_detections(
np.concatenate(np.array(self.detections_history)),
self.frame_shape,
threshold=min(len(self.detections_history), 15)
)
return detections
def detect_vehicles(self, image):
"""
Perform a full window passthrough in the specified frame.
Parameters
----------
image : Current frame.
"""
scales = np.array([.3, .5, .65, .8])
y_top = np.array([.6, .57, .56, .55])
frame_detections = np.empty([0, 4], dtype=np.int64)
for scale, y in zip(scales, y_top):
scale_detections = self.detections_for_scale(image, scale, y, 64)
frame_detections = np.append(frame_detections, scale_detections, axis=0)
detections, self.heatmap = self.merge_detections(frame_detections, image.shape, threshold=1)
self.detections_history.append(detections)
def detections_for_scale(self, image, scale, y, k):
"""
Runs a classifier on all windows for specified frame scale.
Parameters
----------
image : Current frame.
scale : Scale of the image.
y : Top Y coordinate of the windows.
k : Size of the window.
Returns
-------
Boundaries of windows that got detections.
"""
(h, w, d) = image.shape
scaled = resize((image / 255.).astype(np.float64), (int(h * scale), int(w * scale), d), preserve_range=True).astype(np.float32)
extractor = FeatureExtractor(scaled)
(h, w, d) = scaled.shape
detections = np.empty([0, 4], dtype=np.int)
y = int(h*y)
s = k // 3
x_range = np.linspace(0, w - k, (w + s) // s)
for x in x_range.astype(np.int):
features = extractor.feature_vector(x, y, k)
features = self.scaler.transform(np.array(features).reshape(1, -1))
if self.classifier.predict(features)[0] == 1:
detections = np.append(detections, [[x, y, x + k, y + k]], axis=0)
return (detections / scale).astype(np.int)
def add_heat(self, heatmap, coordinates):
"""
Adds a 1 for pixels inside each detected region.
Parameters
----------
heatmap : Array with a heatmap.
coordinates : Detections to merge.
Returns
-------
Updated heatmap.
"""
for c in coordinates:
# Assuming each set of coordinates takes the form (x1, y1, x2, y2)
heatmap[c[1]:c[3], c[0]:c[2]] += 1
return heatmap
def merge_detections(self, detections, image_shape, threshold):
"""
Merges specified detections based on a heatmap and threshold.
Parameters
----------
detections : Array of detections to merge.
image_shape : Shape of the image.
threshold : Heatmap threshold.
Returns
-------
Tuple of merged regions and a heatmap.
"""
heatmap = np.zeros((image_shape[0], image_shape[1])).astype(np.float)
# Add heat to each box in box list
heatmap = self.add_heat(heatmap, detections)
# Apply threshold to help remove false positives
heatmap[heatmap < threshold] = 0
heatmap = np.clip(heatmap, 0, 255)
labels = label(heatmap)
cars = np.empty([0, 4], dtype=np.int64)
# Iterate through all detected cars
for car in range(1, labels[1] + 1):
# Find pixels with each car_number label value
nonzero = (labels[0] == car).nonzero()
cars = np.append(
cars,
[[np.min(nonzero[1]), np.min(nonzero[0]), np.max(nonzero[1]), np.max(nonzero[0])]],
axis=0
)
# Return the image
return (cars, heatmap)