Skip to content

강좌 6 OpenCV를 이용한 차선 검출

Sikurity edited this page Nov 24, 2017 · 1 revision
import numpy as np
from PIL import ImageGrab
import cv2
import time
from numpy import ones,vstack
from numpy.linalg import lstsq
from directkeys import PressKey, W, A, S, D
from statistics import mean

# 4강 참고
def roi(img, vertices):
    
    # 이미지 크기 만큼의 값이 0으로 채워진 배열 생성
    mask = np.zeros_like(img)   
    
    # mask에서 관심영역(vertices)에만 값을 255(=0x11111111)로 변환
    cv2.fillPoly(mask, vertices, 255)
    
    # mask에서 값이 1인 부분의 img만 사용하고 0인 부분은 버림
    masked = cv2.bitwise_and(img, mask)

    return masked

# 차선 그리기, 예외 처리 적용
def draw_lanes(img, lines, color=[0, 255, 255], thickness=3):

    # 에러 발생 시, 기본선을 사용
    try:
        # 차량이 움직이면 화면에서 수평선이 항상 같은 지점에 있지 않게 되므로
        # 차선 마커의 최대 y 값을 찾는다

        ys = []  
        for i in lines:
            for ii in i:
                ys += [ii[1], ii[3]] # 직선의 출발점과 도착점의 y 성분들

        min_y = min(ys)
        max_y = 600

        new_lines = []
        line_dict = {}

        for idx, i in enumerate(lines):
            for xyxy in i:
                # http://stackoverflow.com/questions/21565994/method-to-return-the-equation-of-a-straight-line-given-two-points
                # 두 점의 좌표로 선을 생성
                x_coords = (xyxy[0], xyxy[2])
                y_coords = (xyxy[1], xyxy[3])

                # [ [xyxy[0], xyxy[2]],    
                #   [      1,       1] ].T =
                # [ [xyxy[0], 1],
                #   [xyxy[2], 1] ]
                A = vstack([x_coords, ones(len(x_coords))]).T # 전치 행렬

                m, b = lstsq(A, y_coords)[0]  # y_coords = mA + b의 Least Square(최소제곱해) 반환(feat. 선형대수)

                # y = mx + b 였으므로, x = (y - b) / m
                x1 = (min_y - b) / m
                x2 = (max_y - b) / m

                # 차선 추가
                line_dict[idx] = [m, b, [int(x1), min_y, int(x2), max_y]]
                new_lines.append([int(x1), min_y, int(x2), max_y])

        # ... 여기부터는 제가 영어가 부족해서... 코드 해석 좀 대신 부탁드릴께요!!!
        final_lanes = {}

        for idx in line_dict:
            final_lanes_copy = final_lanes.copy()

            m = line_dict[idx][0]
            b = line_dict[idx][1]
            line = line_dict[idx][2]
            
            if len(final_lanes) == 0:
                final_lanes[m] = [ [m,b,line] ]
                
            else:
                found_copy = False

                for other_ms in final_lanes_copy:

                    if not found_copy:
                        if abs(other_ms*1.2) > abs(m) > abs(other_ms*0.8):
                            if abs(final_lanes_copy[other_ms][0][1]*1.2) > abs(b) > abs(final_lanes_copy[other_ms][0][1]*0.8):
                                final_lanes[other_ms].append([m,b,line])
                                found_copy = True
                                break
                        else:
                            final_lanes[m] = [ [m,b,line] ]

        line_counter = {}

        for lanes in final_lanes:
            line_counter[lanes] = len(final_lanes[lanes])

        top_lanes = sorted(line_counter.items(), key=lambda item: item[1])[::-1][:2]

        lane1_id = top_lanes[0][0]
        lane2_id = top_lanes[1][0]

        # 차선 후보들을 평균 내서, 최종 차선 반환
        def average_lane(lane_data):
            x1s = []
            y1s = []
            x2s = []
            y2s = []
            for data in lane_data:
                x1s.append(data[2][0])
                y1s.append(data[2][1])
                x2s.append(data[2][2])
                y2s.append(data[2][3])

            return int(mean(x1s)), int(mea(y1s)), int(mean(x2s)), int(mean(y2s)) 

        l1_x1, l1_y1, l1_x2, l1_y2 = average_lane(final_lanes[lane1_id])
        l2_x1, l2_y1, l2_x2, l2_y2 = average_lane(final_lanes[lane2_id])

        return [l1_x1, l1_y1, l1_x2, l1_y2], [l2_x1, l2_y1, l2_x2, l2_y2]

    # 예외처리
    except Exception as e:
        print(str(e))

# BGR to Gray -> Canny Edge Detection -> 가우시안 정규화 -> 관심영역 필터링 -> 허프 직선검출 -> 차선 검출 -> 차선 그리기
def process_img(image):
    original_image = image
    # convert to gray
    processed_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # edge detection
    processed_img =  cv2.Canny(processed_img, threshold1 = 200, threshold2=300)
    
    processed_img = cv2.GaussianBlur(processed_img,(5,5),0)
    
    vertices = np.array([[10,500],[10,300],[300,200],[500,200],[800,300],[800,500],
                         ], np.int32)

    processed_img = roi(processed_img, [vertices])

    # 5강 참고
    lines = cv2.HoughLinesP(processed_img, 1,  np.pi/180, 180,   20,         15)
    try:
        l1, l2 = draw_lanes(original_image,lines)
        cv2.line(original_image, (l1[0], l1[1]), (l1[2], l1[3]), [0,255,0], 30)
        cv2.line(original_image, (l2[0], l2[1]), (l2[2], l2[3]), [0,255,0], 30)
    except Exception as e:
        print(str(e))
        pass
    try:
        for coords in lines:
            coords = coords[0]
            try:
                cv2.line(processed_img, (coords[0], coords[1]), (coords[2], coords[3]), [255,0,0], 3)
                
                
            except Exception as e:
                print(str(e))
    except Exception as e:
        pass

    return processed_img, original_image

def main():
    last_time = time.time()
    while True:
        screen =  np.array(ImageGrab.grab(bbox=(0,40,800,640)))

        print('Frame took {} seconds'.format(time.time()-last_time))
        last_time = time.time()

        new_screen, original_image = process_img(screen)

        cv2.imshow('window', new_screen)
        cv2.imshow('window2',cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB))
        
        if cv2.waitKey(25) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break

결과물

lane-finder-example