[ROS] line_drive.py

to move and edit some code

 

directory to be saved : line_drive/

(I think it wasn't run by rosrun or roslaunch or sth. just located in "line_drive" directory with "kmu_track.mkv" and "steer_arrow.png" files )

 

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import rospy
import numpy as np
import cv2, random, math, time

Width = 640
Height = 480
Offset = 340

scan_width_200 = 200
scan_height_20 = 20
lmid = scan_width_200
rmid = Width - scan_width_200
area_width_10 = 10
area_height_10 = 10
row_begin_5 = (scan_height_20 - area_height_10) //2
row_end_15 = row_begin_5 + area_height_10

pixel_threshold = 0.05 * area_width_10 * area_height_10 # 5% threshold

steer_angle_max = 20.0

lpos = -1
rpos = -1


# draw rectangle
def draw_rectangle(img, lpos, rpos, offset=0):

    center = (lpos + rpos) / 2

    if lpos == -1:
        pass
    else:
        cv2.rectangle(img, (lpos - area_width_10//2 , offset + (scan_height_20-area_height_10)//2),
                           (lpos + area_width_10//2 , offset + (scan_height_20-area_height_10)//2 + area_height_10 ),
                           (0, 255, 0), 2)
    #draw green rectangle on a left side line


    if rpos == -1:
        pass
    else:
        cv2.rectangle(img,  (rpos - area_width_10//2, offset + (scan_height_20-area_height_10)//2),
                            (rpos + area_width_10//2, offset + (scan_height_20-area_height_10)//2 + area_height_10),
                            (0, 255, 0), 2)   
    #draw green rectangle on a right side line

    cv2.rectangle(img,      (center - area_width_10//2, offset + (scan_height_20-area_height_10)//2),
                            (center + area_width_10//2, offset + (scan_height_20-area_height_10)//2 + area_height_10),
                            (0,255,0), 2)
    #draw green rectangle on a middle of the lines 

    cv2.rectangle(img,      (Width//2 - area_width_10//2, offset + (scan_height_20-area_height_10)//2),
                            (Width//2 + area_width_10//2, offset + (scan_height_20-area_height_10)//2 + area_height_10),
                            (0,0,255), 2)
    #draw red rectangle on a middle of the camera frame

    return img

# You are to find "left and light position" of road lanes
def process_image(frame):
    global Offset
    global lpos, rpos
    #lpos, rpos : allocate previous value when the line isn't detected

    roi = frame[Offset:Offset+scan_height_20, :]
    #make roi frame

    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    #gray filtering - filter the roi frame into gray image

    kernel_size = 5
    blur = cv2.GaussianBlur(gray, (kernel_size, kernel_size), 0)
    #blur filtering - filter the gray image into blurred image

    low_threshold = 60
    high_threshold = 70
    edge_img = cv2.Canny(np.uint8(blur) , low_threshold, high_threshold)
    #canny filtering - filter the blurred image into canny edge image

    line_search(edge_img)
    #line search algorithm

    frame = draw_rectangle(frame, lpos, rpos, offset=Offset)
    #draw 4 rectangles 
    
    return (lpos, rpos), frame

#line search algorithm
def line_search(edge_img):
    global lpos
    global rpos


    if lpos == -1:
        for l in range((area_width_10//2), lmid - (area_width_10//2)):
            area = edge_img[row_begin_5:row_end_15 , l-(area_width_10//2):l+(area_width_10//2)]
            if cv2.countNonZero(area) > pixel_threshold:
                lpos = l
                break
    #left line initialization
    if rpos == -1:
        for r in range(Width - (area_width_10//2), rmid + (area_width_10//2) , -1):
            area = edge_img[row_begin_5:row_end_15 , r-(area_width_10//2):r+(area_width_10//2)]
            if cv2.countNonZero(area) > pixel_threshold:
                rpos = r
                break
    #right line initialization


    if lpos < 5: 
        lpos = 5
    if rpos > Width - 5:
        rpos = Width - 5
    #lpos, rpos adjustment , when it's gone out of camera frame (for avoiding detecting middle line)


    l = lpos
    l_step = 1
    search_limit = Width//7
    while( l_step < search_limit):
        area = edge_img[row_begin_5:row_end_15 , l-(area_width_10//2):l+(area_width_10//2)]
        if cv2.countNonZero(area) > pixel_threshold:
            lpos = l
            break
        l = l + ((-1)**(l_step+1)) *l_step
        l_step = l_step + 1
    #left line searching - search from the last location of left line
    r = rpos
    r_step = 1
    search_limit = Width//7
    while( r_step < search_limit):
        area = edge_img[row_begin_5:row_end_15 , r-(area_width_10//2):r+(area_width_10//2)]
        if cv2.countNonZero(area) > pixel_threshold:
            rpos = r
            break
        r = r + ((-1)**(r_step+1)) *r_step
        r_step = r_step + 1
    #right line searching - search from the last location of right line


def draw_steer(image, steer_angle):
    global Width, Height, arrow_pic

    arrow_pic = cv2.imread('steer_arrow.png', cv2.IMREAD_COLOR)

    origin_Height = arrow_pic.shape[0]
    origin_Width = arrow_pic.shape[1]
    steer_wheel_center = origin_Height * 0.74
    arrow_Height = Height/2
    arrow_Width = (arrow_Height * 462)/728

    matrix = cv2.getRotationMatrix2D((origin_Width/2, steer_wheel_center), (steer_angle) * 1.5, 0.7)    
    arrow_pic = cv2.warpAffine(arrow_pic, matrix, (origin_Width+60, origin_Height))
    arrow_pic = cv2.resize(arrow_pic, dsize=(arrow_Width, arrow_Height), interpolation=cv2.INTER_AREA)

    gray_arrow = cv2.cvtColor(arrow_pic, cv2.COLOR_BGR2GRAY)
    _, mask = cv2.threshold(gray_arrow, 1, 255, cv2.THRESH_BINARY_INV)

    arrow_roi = image[arrow_Height: Height, (Width/2 - arrow_Width/2) : (Width/2 + arrow_Width/2)]
    arrow_roi = cv2.add(arrow_pic, arrow_roi, mask=mask)
    res = cv2.add(arrow_roi, arrow_pic)
    image[(Height - arrow_Height): Height, (Width/2 - arrow_Width/2): (Width/2 + arrow_Width/2)] = res

    cv2.imshow('steer', image)


# You are to publish "steer_anlge" following load lanes
if __name__ == '__main__':
    cap = cv2.VideoCapture('kmu_track.mkv')
    time.sleep(3)

    while not rospy.is_shutdown():
        ret, image = cap.read()
        pos, frame = process_image(image)
        
        shift = float(Width//2 - (pos[0]+pos[1])/2 ) #shift = distance from the absolute middle of the frame to the middle of left and right lines
        shift_max = Width//10


        if shift >=0.0:
            steer_angle = ( (steer_angle_max)/((shift_max)**2) ) * (shift**2)
        else :
            steer_angle = - ( (steer_angle_max)/((shift_max)**2) ) * (shift**2)
        if steer_angle>steer_angle_max:
            steer_angle=steer_angle_max
        if steer_angle<-steer_angle_max:
            steer_angle=-steer_angle_max

        
        draw_steer(frame, steer_angle)

        if cv2.waitKey(3) & 0xFF == ord('q'):
            break