Permalink
Cannot retrieve contributors at this time
import numpy as np | |
import cv2 | |
from lanetracker.window import Window | |
from lanetracker.line import Line | |
from lanetracker.gradients import get_edges | |
from lanetracker.perspective import flatten_perspective | |
class LaneTracker(object): | |
""" | |
Tracks the lane in a series of consecutive frames. | |
""" | |
def __init__(self, first_frame, n_windows=9): | |
""" | |
Initialises a tracker object. | |
Parameters | |
---------- | |
first_frame : First frame of the frame series. We use it to get dimensions and initialise values. | |
n_windows : Number of windows we use to track each lane edge. | |
""" | |
(self.h, self.w, _) = first_frame.shape | |
self.win_n = n_windows | |
self.left = None | |
self.right = None | |
self.l_windows = [] | |
self.r_windows = [] | |
self.initialize_lines(first_frame) | |
def initialize_lines(self, frame): | |
""" | |
Finds starting points for left and right lines (e.g. lane edges) and initialises Window and Line objects. | |
Parameters | |
---------- | |
frame : Frame to scan for lane edges. | |
""" | |
# Take a histogram of the bottom half of the image | |
edges = get_edges(frame) | |
(flat_edges, _) = flatten_perspective(edges) | |
histogram = np.sum(flat_edges[int(self.h / 2):, :], axis=0) | |
nonzero = flat_edges.nonzero() | |
# Create empty lists to receive left and right lane pixel indices | |
l_indices = np.empty([0], dtype=np.int) | |
r_indices = np.empty([0], dtype=np.int) | |
window_height = int(self.h / self.win_n) | |
for i in range(self.win_n): | |
l_window = Window( | |
y1=self.h - (i + 1) * window_height, | |
y2=self.h - i * window_height, | |
x=self.l_windows[-1].x if len(self.l_windows) > 0 else np.argmax(histogram[:self.w // 2]) | |
) | |
r_window = Window( | |
y1=self.h - (i + 1) * window_height, | |
y2=self.h - i * window_height, | |
x=self.r_windows[-1].x if len(self.r_windows) > 0 else np.argmax(histogram[self.w // 2:]) + self.w // 2 | |
) | |
# Append nonzero indices in the window boundary to the lists | |
l_indices = np.append(l_indices, l_window.pixels_in(nonzero), axis=0) | |
r_indices = np.append(r_indices, r_window.pixels_in(nonzero), axis=0) | |
self.l_windows.append(l_window) | |
self.r_windows.append(r_window) | |
self.left = Line(x=nonzero[1][l_indices], y=nonzero[0][l_indices], h=self.h, w = self.w) | |
self.right = Line(x=nonzero[1][r_indices], y=nonzero[0][r_indices], h=self.h, w = self.w) | |
def scan_frame_with_windows(self, frame, windows): | |
""" | |
Scans a frame using initialised windows in an attempt to track the lane edges. | |
Parameters | |
---------- | |
frame : New frame | |
windows : Array of windows to use for scanning the frame. | |
Returns | |
------- | |
A tuple of arrays containing coordinates of points found in the specified windows. | |
""" | |
indices = np.empty([0], dtype=np.int) | |
nonzero = frame.nonzero() | |
window_x = None | |
for window in windows: | |
indices = np.append(indices, window.pixels_in(nonzero, window_x), axis=0) | |
window_x = window.mean_x | |
return (nonzero[1][indices], nonzero[0][indices]) | |
def process(self, frame, draw_lane=True, draw_statistics=True): | |
""" | |
Performs a full lane tracking pipeline on a frame. | |
Parameters | |
---------- | |
frame : New frame to process. | |
draw_lane : Flag indicating if we need to draw the lane on top of the frame. | |
draw_statistics : Flag indicating if we need to render the debug information on top of the frame. | |
Returns | |
------- | |
Resulting frame. | |
""" | |
edges = get_edges(frame) | |
(flat_edges, unwarp_matrix) = flatten_perspective(edges) | |
(l_x, l_y) = self.scan_frame_with_windows(flat_edges, self.l_windows) | |
self.left.process_points(l_x, l_y) | |
(r_x, r_y) = self.scan_frame_with_windows(flat_edges, self.r_windows) | |
self.right.process_points(r_x, r_y) | |
if draw_statistics: | |
edges = get_edges(frame, separate_channels=True) | |
debug_overlay = self.draw_debug_overlay(flatten_perspective(edges)[0]) | |
top_overlay = self.draw_lane_overlay(flatten_perspective(frame)[0]) | |
debug_overlay = cv2.resize(debug_overlay, (0, 0), fx=0.3, fy=0.3) | |
top_overlay = cv2.resize(top_overlay, (0, 0), fx=0.3, fy=0.3) | |
frame[:250, :, :] = frame[:250, :, :] * .4 | |
(h, w, _) = debug_overlay.shape | |
frame[20:20 + h, 20:20 + w, :] = debug_overlay | |
frame[20:20 + h, 20 + 20 + w:20 + 20 + w + w, :] = top_overlay | |
text_x = 20 + 20 + w + w + 20 | |
self.draw_text(frame, 'Radius of curvature: {} m'.format(self.radius_of_curvature()), text_x, 80) | |
self.draw_text(frame, 'Distance (left): {:.1f} m'.format(self.left.camera_distance()), text_x, 140) | |
self.draw_text(frame, 'Distance (right): {:.1f} m'.format(self.right.camera_distance()), text_x, 200) | |
if draw_lane: | |
frame = self.draw_lane_overlay(frame, unwarp_matrix) | |
return frame | |
def draw_text(self, frame, text, x, y): | |
cv2.putText(frame, text, (x, y), cv2.FONT_HERSHEY_SIMPLEX, .8, (255, 255, 255), 2) | |
def draw_debug_overlay(self, binary, lines=True, windows=True): | |
""" | |
Draws an overlay with debugging information on a bird's-eye view of the road (e.g. after applying perspective | |
transform). | |
Parameters | |
---------- | |
binary : Frame to overlay. | |
lines : Flag indicating if we need to draw lines. | |
windows : Flag indicating if we need to draw windows. | |
Returns | |
------- | |
Frame with an debug information overlay. | |
""" | |
if len(binary.shape) == 2: | |
image = np.dstack((binary, binary, binary)) | |
else: | |
image = binary | |
if windows: | |
for window in self.l_windows: | |
coordinates = window.coordinates() | |
cv2.rectangle(image, coordinates[0], coordinates[1], (1., 1., 0), 2) | |
for window in self.r_windows: | |
coordinates = window.coordinates() | |
cv2.rectangle(image, coordinates[0], coordinates[1], (1., 1., 0), 2) | |
if lines: | |
cv2.polylines(image, [self.left.get_points()], False, (1., 0, 0), 2) | |
cv2.polylines(image, [self.right.get_points()], False, (1., 0, 0), 2) | |
return image * 255 | |
def draw_lane_overlay(self, image, unwarp_matrix=None): | |
""" | |
Draws an overlay with tracked lane applying perspective unwarp to project it on the original frame. | |
Parameters | |
---------- | |
image : Original frame. | |
unwarp_matrix : Transformation matrix to unwarp the bird's eye view to initial frame. Defaults to `None` (in | |
which case no unwarping is applied). | |
Returns | |
------- | |
Frame with a lane overlay. | |
""" | |
# Create an image to draw the lines on | |
overlay = np.zeros_like(image).astype(np.uint8) | |
points = np.vstack((self.left.get_points(), np.flipud(self.right.get_points()))) | |
# Draw the lane onto the warped blank image | |
cv2.fillPoly(overlay, [points], (0, 255, 0)) | |
if unwarp_matrix is not None: | |
# Warp the blank back to original image space using inverse perspective matrix (Minv) | |
overlay = cv2.warpPerspective(overlay, unwarp_matrix, (image.shape[1], image.shape[0])) | |
# Combine the result with the original image | |
return cv2.addWeighted(image, 1, overlay, 0.3, 0) | |
def radius_of_curvature(self): | |
""" | |
Calculates radius of the lane curvature by averaging curvature of the edge lines. | |
Returns | |
------- | |
Radius of the lane curvature in meters. | |
""" | |
return int(np.average([self.left.radius_of_curvature(), self.right.radius_of_curvature()])) |