Commit 3165d0e6 authored by Sparkf's avatar Sparkf 🏙️
Browse files

split the program into module

parent 78c6cbd4

train1/ld_module_binary.py

0 → 100644
+105 −0
Original line number Diff line number Diff line 
import numpy as np

import cv2

import ld_module_transform





def get_bin_img(img, kernel_size=3, sobel_dirn='X', sobel_thresh=(0, 255), r_thresh=(0, 255),

                s_thresh=(0, 255), b_thresh=(0, 255), g_thresh=(0, 255)):

    hls = cv2.cvtColor(img, cv2.COLOR_BGR2HLS).astype(np.float32)

    h_channel = hls[:, :, 0]

    l_channel = hls[:, :, 1]

    s_channel = hls[:, :, 2]



    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)



    if sobel_dirn == 'X':

        sobel = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=kernel_size)

    else:

        sobel = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=kernel_size)



    abs_sobel = np.absolute(sobel)

    scaled_sobel = np.uint8(255 * abs_sobel / np.max(abs_sobel))



    sbinary = np.zeros_like(scaled_sobel)

    sbinary[(scaled_sobel >= sobel_thresh[0]) & (scaled_sobel <= sobel_thresh[1])] = 1



    combined = np.zeros_like(sbinary)

    combined[(sbinary == 1)] = 1



    # Threshold R color channel

    r_binary = get_rgb_thresh_img(img, thresh=r_thresh)



    # Threshhold G color channel

    g_binary = get_rgb_thresh_img(img, thresh=g_thresh, channel='G')



    # Threshhold B in LAB

    b_binary = get_lab_bthresh_img(img, thresh=b_thresh)



    # Threshold color channel

    s_binary = get_hls_sthresh_img(img, thresh=s_thresh)



    # If two of the three are activated, activate in the binary image

    combined_binary = np.zeros_like(combined)

    combined_binary[(r_binary == 1) | (combined == 1) | (s_binary == 1) | (b_binary == 1) | (g_binary == 1)] = 1



    return combined_binary





def get_hls_sthresh_img(img, thresh=(0, 255)):

    hls_img = cv2.cvtColor(img, cv2.COLOR_BGR2HLS)

    S = hls_img[:, :, 2]



    binary_output = np.zeros_like(S).astype(np.uint8)

    binary_output[(S >= thresh[0]) & (S < thresh[1])] = 1



    return binary_output





def get_lab_bthresh_img(img, thresh=(0, 255)):

    lab_img = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)

    B = lab_img[:, :, 2]



    bin_op = np.zeros_like(B).astype(np.uint8)

    bin_op[(B >= thresh[0]) & (B < thresh[1])] = 1



    return bin_op





def get_rgb_thresh_img(img, channel='R', thresh=(0, 255)):

    img1 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    if channel == 'R':

        bin_img = img1[:, :, 0]

    if channel == 'G':

        bin_img = img1[:, :, 1]

    if channel == 'B':

        bin_img = img1[:, :, 2]



    binary_img = np.zeros_like(bin_img).astype(np.uint8)

    binary_img[(bin_img >= thresh[0]) & (bin_img < thresh[1])] = 1



    return binary_img





def abs_thresh(img, sobel_kernel=9, mag_thresh=(0, 255), return_grad=False, direction='y'):

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    grad = None

    scaled_sobel = None



    # Sobel x

    if direction.lower() == 'x':

        grad = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=sobel_kernel)  # Take the derivative in x

    # Sobel y

    else:

        grad = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=sobel_kernel)  # Take the derivative in y



    if return_grad == True:

        return grad



    abs_sobel = np.absolute(grad)  # Absolute x derivative to accentuate lines away from horizontal

    scaled_sobel = np.uint8(255 * abs_sobel / np.max(abs_sobel))



    grad_binary = np.zeros_like(scaled_sobel)

    grad_binary[(scaled_sobel >= mag_thresh[0]) & (scaled_sobel < mag_thresh[1])] = 1



    return grad_binary

train1/ld_module_detection.py

0 → 100644
+89 −0
Original line number Diff line number Diff line 
import numpy as np

import cv2

from scipy.signal import find_peaks

from scipy.optimize import curve_fit

import matplotlib.pyplot as plt

import circle_fit

import os

import json

from skimage.filters import threshold_otsu

import time





def intensity_map(input_array, distance):

    histogram = np.sum(input_array[input_array.shape[0] // 2:, :], axis=0)

    peak_array = find_peaks(histogram, distance=distance)[0]

    return peak_array





def find_point(input_img, start_W):



    midpoint_arr_R = a = np.zeros((0, 2)).astype(int)

    # print(midpoint_arr_R)

    midpoint_tmp_R, bottom_R = find_hist_iteration(input_img, 450, start_W)

    # print(midpoint_tmp_R)



    while bottom_R > 0:

        window_H, window_WM, find_indicator = gen_windows(input_img,bottom_R, midpoint_tmp_R, 15, 100)

        midpoint_tmp_R = int(window_WM)





        if find_indicator == 1:

            tmp_list = np.array([[midpoint_tmp_R, bottom_R]]).astype(int)

            midpoint_arr_R = np.vstack([midpoint_arr_R, tmp_list])

            # plt.scatter(midpoint_tmp_R, bottom_R)

        bottom_R = bottom_R - 15

    return midpoint_arr_R





def find_hist_iteration(input_img, baseline=500, start_W=0):

    input_img_hist_iter = input_img[baseline:baseline + 149, start_W:start_W + 200]

    find_hist_iteration_iter = np.sum(input_img_hist_iter[input_img_hist_iter.shape[0] // 2:, :], axis=0)

    peak_hist = find_peaks(find_hist_iteration_iter, distance=250)[0] or [0]

    if (find_hist_iteration_iter[peak_hist[0]] < 20) or peak_hist.size == 0:

        find_hist_iteration(baseline - 50, start_W)

    else:

        pass

    return (peak_hist[0] + start_W), baseline + 50





def gen_windows(input_img, bottom_H, midpoint_W, H, W):

    find_indicator = 1

    window = input_img[bottom_H - 15:bottom_H, midpoint_W - 50:midpoint_W + 50]

    H = 15

    W = 100

    # window_arr = input_img[]

    histogram_window = np.sum(window[window.shape[0] // 2:, :], axis=0)

    window_peak = find_peaks(histogram_window, distance=100)[0]



    return_mid = 0



    if window_peak.size > 0:

        if histogram_window[window_peak] > 5 and np.abs(window_peak - (W / 2)) < 40:

            return_mid = midpoint_W - (W / 2) + window_peak

        else:

            return_mid = midpoint_W

    else:

        return_mid = midpoint_W



        find_indicator = 0

    # plt.plot(histogram_window)

    # plt.show()

    return bottom_H - H / 2, return_mid, find_indicator



def fit_lane(input_img,midpoint):

    midpoint_arr = np.zeros((0, 2)).astype(int)

    xcl = 0

    ycl = 0

    rcl = 0



    try:

        midpoint_arr = find_point(input_img,midpoint - 100)

        # plt.scatter(x=midpoint_arr.T[0], y=midpoint_arr.T[1], s=8)

        ycl, xcl, rcl, _ = circle_fit.hyper_fit(midpoint_arr)

        # print(rcl)

        # if (rcl > 1000):

        #     plt.gca().add_artist(plt.Circle((ycl, xcl), rcl, color='r', fill=False))

    except Exception:

        pass

    return midpoint_arr, ycl, xcl, rcl
 
 No newline at end of file

train1/ld_module_transform.py

0 → 100644
+24 −0
Original line number Diff line number Diff line 
import numpy as np

import cv2

from scipy.signal import find_peaks

from scipy.optimize import curve_fit

import matplotlib.pyplot as plt

import circle_fit

import os

import json

from skimage.filters import threshold_otsu

import time

import ld_module_binary







def transform_to(input_img):



    pts1 = np.float32([[450, 300], [-851, 720], [870, 300], [2057, 720]])

    pts2 = np.float32([[0, 0], [0, 600], [600, 0], [600, 600]])



    M = cv2.getPerspectiveTransform(pts1, pts2)



    dst = cv2.warpPerspective(input_img, M, (600, 600))



    return dst
 
 No newline at end of file

train1/main.py

0 → 100644
+117 −0
Original line number Diff line number Diff line 
import cv2

import os

import time

from numba import njit, prange



import matplotlib.pyplot as plt



import ld_module_transform

import ld_module_binary

import numpy as np

import ld_module_detection



start_time = time.time()



# rootdir = 'D:/train_set/clips/0531/'

rootdir = 'D:/train_set/clips/examine/'

foldername = "null"







for subdir, dirs, files in os.walk(rootdir):

    for d in dirs:

        # for file in files:

        #     impath = os.path.join(subdir, file)

        # read into seq

        # buffer first three image



        # img_buffer_array = np.zeros([4, 720, 1280, 3]).astype('uint8')

        # for k1 in range(1, 4):

        #     print(subdir + d + "/" + str(k1) + ".jpg")

        #     input_img = cv2.imread(subdir + d + "/" + str(k1) + ".jpg")

        #     input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)

        #     img_buffer_array[k1 - 1] = input_img

        #     print(k1)

        # rotate_num = 3



        for k2 in range(1, 21):

            input_img = cv2.imread(subdir + d + "/" + str(k2) + ".jpg")

            input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)



            # yellow filter

            input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)

            original_img = input_img

            # img_buffer_array[rotate_num] = input_img

            # input_img = (img_buffer_array[(rotate_num) % 4] * 0.25 + img_buffer_array[(rotate_num - 1) % 4] * 0.25 +

            #              img_buffer_array[

            #                  (rotate_num - 2) % 4] * 0.25 + img_buffer_array[(rotate_num - 3) % 4] * 0.25).astype(

            #     'uint8')

            # input_img = np.maximum.reduce([img_buffer_array[(rotate_num) % 4], img_buffer_array[(rotate_num - 1) % 4],

            #                                img_buffer_array[(rotate_num - 2) % 4],

            #                                img_buffer_array[(rotate_num - 3) % 4]]).astype('uint8')

            # input_img = input_img[:,:,0]

            # foldername = subdir[-19:]

            foldername = d

            filename = str(k2)

            i = 0



            # 1 perspective transformation

            transformed_img = ld_module_transform.transform_to(input_img)



            # 2 input for binary operation

            kernel_size = 7

            mag_thresh = (30, 150)

            r_thresh = (235, 255)

            s_thresh = (165, 255)

            b_thresh = (160, 255)

            g_thresh = (210, 255)

            extracted_img = ld_module_binary.get_bin_img(transformed_img, kernel_size=kernel_size,

                                                         sobel_thresh=mag_thresh,

                                                         r_thresh=r_thresh,

                                                         s_thresh=s_thresh, b_thresh=b_thresh,

                                                         g_thresh=g_thresh) \

                            + ld_module_binary.abs_thresh(transformed_img,

                                                          sobel_kernel=3,

                                                          mag_thresh=(

                                                              35, 210),

                                                          direction='x')



            # 3 add line mark

            # 3.1 determine the start point of the mark

            # 3.1.1 eliminate the lower L/R traingle

            input_img = extracted_img

            tri_matrix = np.tri(600, 600, 420).T

            input_img = input_img * tri_matrix * np.flip(np.tri(600, 600, 440).T, 1)



            # 3.1.2 find peak (line) number

            img_peak_histogram = ld_module_detection.intensity_map(input_img, 150)



            # 3.1.3 find lane via sliding window

            for midpoint in img_peak_histogram:

                # midpoint (width)

                mark_layer = np.zeros((600, 600, 3), np.uint8)



                if 100 < midpoint < 500:

                    point_array, ycl, xcl, rcl = ld_module_detection.fit_lane(input_img, midpoint)



                    # 3.1.3.1 draw line over the original image

                    if point_array.size > 0 and rcl > 800 and rcl < 2147483647:

                        for point in point_array:

                            cv2.circle(mark_layer, tuple(point), 10, (0, 0, 255), thickness=-1)

                            cv2.circle(mark_layer, tuple([int(ycl), int(xcl)]), int(rcl), (0, 255, 0), thickness=10)



                # 3.1.3.2 transform the mark layer back

                pts2 = np.float32([[450, 300], [-851, 720], [870, 300], [2057, 720]])

                pts1 = np.float32([[0, 0], [0, 600], [600, 0], [600, 600]])



                M1 = cv2.getPerspectiveTransform(pts1, pts2)



                mark_layer_trans = cv2.warpPerspective(mark_layer, M1, (1280, 720))

                mark_layer_mask = np.where(mark_layer_trans > 0, 0, 1)

                original_img = original_img * mark_layer_mask + mark_layer_trans



                # 4 save image

                cv2.imwrite("op_module/" + foldername + "_" + filename + ".png", original_img)





print("--- %s seconds ---" % (time.time() - start_time))

train1/pre-processing.py

+22 −21
Original line number Diff line number Diff line
@@ -285,6 +285,7 @@ def add_line_mark(input_img, dst, original_img): 


    mark_layer = np.zeros((600, 600, 3), np.uint8)

    for j in peak_hist_all[0]:

        if 100 < j < 500:

            point_array,ycl,xcl,rcl = fit_lane(j)



            if point_array.size > 0 and rcl > 800 and rcl < 2147483647:
@@ -301,7 +302,7 @@ def add_line_mark(input_img, dst, original_img): 
    mark_layer_trans = cv2.warpPerspective(mark_layer, M1, (1280,720))

    mark_layer_mask = np.where(mark_layer_trans > 0,0,1)

    original_img = original_img*mark_layer_mask + mark_layer_trans

    plt.imshow(original_img)



    # mark_layer_trans = cv2.addWeighted(mark_layer_trans, 1, original_img, 1, 0.0)
@@ -375,27 +376,27 @@ for subdir, dirs, files in os.walk(rootdir): 
        # read into seq

        # buffer first three image



        img_buffer_array = np.zeros([4, 720, 1280, 3]).astype('uint8')

        for k1 in range(1, 4):

            print(subdir + d + "/" + str(k1) + ".jpg")

            input_img = cv2.imread(subdir + d + "/" + str(k1) + ".jpg")

            input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)

            img_buffer_array[k1 - 1] = input_img

            print(k1)

        rotate_num = 3

        # img_buffer_array = np.zeros([4, 720, 1280, 3]).astype('uint8')

        # for k1 in range(1, 4):

        #     print(subdir + d + "/" + str(k1) + ".jpg")

        #     input_img = cv2.imread(subdir + d + "/" + str(k1) + ".jpg")

        #     input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)

        #     img_buffer_array[k1 - 1] = input_img

        #     print(k1)

        # rotate_num = 3



        for k2 in range(4, 21):

        for k2 in range(1, 21):

            input_img = cv2.imread(subdir + d + "/" + str(k2) + ".jpg")

            input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)



            # yellow filter

            input_img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB)

            original_img = input_img

            img_buffer_array[rotate_num] = input_img

            input_img = (img_buffer_array[(rotate_num) % 4] * 0.25 + img_buffer_array[(rotate_num - 1) % 4] * 0.25 +

                         img_buffer_array[

                             (rotate_num - 2) % 4] * 0.25 + img_buffer_array[(rotate_num - 3) % 4] * 0.25).astype(

                'uint8')

            # img_buffer_array[rotate_num] = input_img

            # input_img = (img_buffer_array[(rotate_num) % 4] * 0.25 + img_buffer_array[(rotate_num - 1) % 4] * 0.25 +

            #              img_buffer_array[

            #                  (rotate_num - 2) % 4] * 0.25 + img_buffer_array[(rotate_num - 3) % 4] * 0.25).astype(

            #     'uint8')

            # input_img = np.maximum.reduce([img_buffer_array[(rotate_num) % 4], img_buffer_array[(rotate_num - 1) % 4],

            #                                img_buffer_array[(rotate_num - 2) % 4],

            #                                img_buffer_array[(rotate_num - 3) % 4]]).astype('uint8')

CRA Git | Maintained and supported by SUSTech CRA and CCSE