レーンエッジの取得

4293 ワード

import cv2
import numpy as np

# from numpy.core.fromnumeric import mean

file_name = './lane_detect_test_data1.mp4'
cap = cv2.VideoCapture(file_name)

hue_upper = 179
hue_lower = 0
satu_upper = 255
satu_lower = 0
val_upper = 255
val_lower = 238


def gray_scale(img):
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    return img_gray


def gaussian_blur(img_gray, kernel_size):
    return cv2.GaussianBlur(img_gray, (kernel_size, kernel_size), 0)


def canny(img_gray):
    img_canny = cv2.Canny(img_gray, 0, 70)
    return img_canny


def make_mask(img):
    height, width = img.shape[:2]
    mask = np.zeros_like(img)
    vertices = np.array([[(0, height - 100), (width / 2 - 270, height / 2 + 100), (width / 2 + 250, height / 2 + 100),
                          (width-180, height - 100)]], dtype=np.int32)
    cv2.fillPoly(mask, vertices, (255, 255, 255))
    #cv2.imshow('mask', mask)
    return mask


def find_roi(img, img_canny, mask):
    img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    mask2 = cv2.inRange(img_hsv, (hue_lower, satu_lower, val_lower), (hue_upper, satu_upper, val_upper))
    mask2 = cv2.cvtColor(mask2, cv2.COLOR_GRAY2BGR)
    road_mask = cv2.bitwise_and(img, mask2)
    road_mask = cv2.bitwise_not(road_mask)

    roi = cv2.bitwise_and(road_mask, mask)
    roi = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)

    return roi


def divide_img(roi):
    _, width = roi.shape[:2]

    # 왼쪽 이미지
    img_left = roi[:, 0:width // 2].copy()

    # 오른쪽 이미지
    img_right = roi.copy()
    img_right[:, 0:width // 2] = 0

    return img_right, img_left


def find_rho_theta(roi):
    lines = cv2.HoughLines(roi, 1, np.pi / 180, 73)
    line_rho = []
    line_theta = []

    for i in range(len(lines)):
        for rho, theta in lines[i]:
            line_rho.append(rho)
            line_theta.append(theta)

    rho = np.asarray(line_rho).mean()
    theta = np.asarray(line_theta).mean()
    return rho, theta


def draw_line(img, rho, theta):
    a = np.cos(theta)
    b = np.sin(theta)
    x0 = a * rho
    y0 = b * rho
    x1 = int(x0 + 2000 * (-b))
    y1 = int(y0 + 2000 * a)
    x2 = int(x0 - 2000 * (-b))
    y2 = int(y0 - 2000 * a)
    cv2.line(img, (x1, y1), (x2, y2), (255, 0, 0), 2)

    return img



# 동영상 프레임에 대해 처리
frame_cnt = 0
while cap.isOpened():
    ret, frame = cap.read()

    if ret:
        frame_cnt += 1
        img_gray = gray_scale(frame)
        img_blur = gaussian_blur(img_gray, 7)
        img_canny = canny(img_blur)
        height, width = frame.shape[:2]

        mask = make_mask(img_gray)
        roi = cv2.bitwise_and(img_canny, mask)
        #cv2.imshow('roi', roi)
        divide = []
        divide = divide_img(roi)
        img_right = divide[0]
        img_left = divide[1]

        rho_theta_right = []
        rho_theta_right = find_rho_theta(img_right)
        rho_theta_left = []
        rho_theta_left = find_rho_theta(img_left)
        result = draw_line(frame, rho_theta_right[0], rho_theta_right[1])
        result = draw_line(result, rho_theta_left[0], rho_theta_left[1])
        print(frame_cnt, rho_theta_right[1], rho_theta_left[1])

        cv2.imshow('result', result)

    # q 누르면 중단
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

'''
# image 디버깅
cap.set(cv2.CAP_PROP_POS_FRAMES, 1076)
ret, frame = cap.read()
img_gray = gray_scale(frame)
img_blur = gaussian_blur(img_gray, 7)
img_canny = canny(img_blur)
height, width = frame.shape[:2]

mask = make_mask(img_gray)
roi = cv2.bitwise_and(img_canny, mask)
cv2.imshow('roi', roi)
cv2.imshow('canny', img_canny)
#cv2.imshow('img', frame)
cv2.waitKey(0)

divide = []
divide = divide_img(roi)
img_right = divide[0]
img_left = divide[1]
#cv2.imshow('r', img_right)
# cv2.imshow('l', img_left)
# cv2.waitKey(0)
'''
rho_theta_right = []
rho_theta_right = find_rho_theta(img_right)
rho_theta_left = []
rho_theta_left = find_rho_theta(img_left)
result = draw_line(frame, rho_theta_right[0], rho_theta_right[1])
result = draw_line(result, rho_theta_left[0], rho_theta_left[1])
print(rho_theta_right[1], rho_theta_left[1])

cv2.imshow('result', result)
cv2.waitKey(0)
cap.release()
cv2.destroyAllWindows()