Last active
February 25, 2019 20:45
-
-
Save codebudo/551a52dbc5a05591e70f9b715faef5da to your computer and use it in GitHub Desktop.
This file contains hidden or 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
from picamera.array import PiRGBArray | |
from picamera import PiCamera | |
import time | |
import cv2 | |
import numpy as np | |
import pickle | |
RESOLUTION = (1280, 720) | |
def apply_canny(image): | |
# TODO should a color channel detect orange cones better? | |
gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) | |
blur = cv2.GaussianBlur(gray, (5,5), 0) | |
canny = cv2.Canny(blur, 50, 150) | |
return canny | |
def region_of_interest(image): | |
height = image.shape[0] | |
width = image.shape[1] | |
triangle = [np.array([ | |
(100, height), | |
(width-100, height), | |
(int(width/2), int(height/4)) | |
])] | |
mask = np.zeros_like(image) | |
cv2.fillPoly(mask, triangle, 255) | |
masked_image = cv2.bitwise_and(image, mask) | |
return masked_image | |
def display_lines(image, lines): | |
line_image = np.zeros_like(image) | |
if lines is not None: | |
for line in lines: | |
if line.shape[0] == 4: | |
x1, y1, x2, y2 = line.reshape(4) | |
cv2.line(line_image, (x1, y1), (x2, y2), (255, 0, 0), 10) | |
return line_image | |
def make_coordinates(image, line_parameters): | |
slope, intercept = line_parameters | |
y1 = image.shape[0] | |
y2 = int(y1 * 3/5) | |
x1 = int((y1 - intercept)/ slope) | |
x2 = int((y2 - intercept)/ slope) | |
return np.array([x1, y1, x2, y2]) | |
def average_slope_intercept(image, lines): | |
left_fit = [] | |
right_fit = [] | |
if lines is None: | |
return np.array([np.array([0,0,0,0]), np.array([0,0,0,0])]) | |
for line in lines: | |
x1, y1, x2, y2 = line.reshape(4) | |
parameters = np.polyfit((x1, x2), (y1, y2), 1) | |
slope = parameters[0] | |
intercept = parameters[1] | |
# ignore horizontal lines in slope range between -0.1 to 0.1 | |
if slope < -0.1: | |
left_fit.append((slope, intercept)) | |
elif slope > 0.1: | |
right_fit.append((slope, intercept)) | |
left_fit_average = np.average(left_fit, axis=0) | |
right_fit_average = np.average(right_fit, axis=0) | |
# Only convert to coordinates if we have line data | |
left_line = np.array([]) | |
if len(left_fit_average.shape) > 0: | |
left_line = make_coordinates(image, left_fit_average) | |
right_line = np.array([]) | |
if len(right_fit_average.shape) > 0: | |
right_line = make_coordinates(image, right_fit_average) | |
return np.array([left_line, right_line]) | |
# camera calibration and undistortion | |
fileHandle = open("calibration.pickle",'rb') | |
ret, mtx, dist, rvecs, tvecs = pickle.load(fileHandle) | |
newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, RESOLUTION, 1, RESOLUTION) | |
# initialize the camera and grab a reference to the raw camera capture | |
time.sleep(1) | |
camera = PiCamera() | |
camera.resolution = RESOLUTION | |
camera.framerate = 10 | |
rawCapture = PiRGBArray(camera, size=RESOLUTION) | |
# capture frames from the camera | |
frame_num = 0 | |
for frame in camera.capture_continuous(rawCapture, format="bgr", use_video_port=True): | |
# grab the raw NumPy array representing the image | |
image = frame.array | |
original_image = np.copy(image) | |
# undistort | |
image = cv2.undistort(image, mtx, dist, None, newcameramtx) | |
x, y, w, h = roi | |
image = image[y:y+h, x:x+w] | |
# modify image frame | |
canny = apply_canny(image) | |
# TODO determine what cropped area should be used based on camera placement | |
cropped_image = region_of_interest(canny) | |
# Find lines | |
lines = cv2.HoughLinesP(cropped_image, 2, np.pi/180, 100, np.array([]), | |
minLineLength=40, maxLineGap=5) | |
if lines is not None: | |
averaged_lines = average_slope_intercept(original_image, lines) | |
line_image = display_lines(canny, averaged_lines) | |
# convert line_image to three channel image | |
line_image = cv2.cvtColor(line_image, cv2.COLOR_GRAY2RGB) | |
combo_image = cv2.addWeighted(image, 0.8, line_image, 1, 1) | |
# show the frame | |
#cv2.imshow("Frame", image) | |
#key = cv2.waitKey(1) | |
# only run 10 frames | |
if frame_num < 10: | |
cv2.imwrite('output'+str(frame_num)+'.jpg', combo_image) | |
else: | |
break | |
frame_num += 1 | |
# clear the stream in preparation for the next frame | |
rawCapture.truncate(0) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment