Играем в GTA V c Python. Часть IV: поиск дорожной разметки

Предыдущая статья — Играем в GTA V c Python. Часть III: прямой вход в игру.

Теперь, когда мы можем считывать фреймы и производить прямой ввод команд с клавиатуры, давайте вернемся к нашей основной задаче — беспилотному вождению автомобиля.

Начнем мы с определения полос движения для нашего транспортного средства. Мы можем использовать обнаружение полос движения как для создания автономного AI, который работает на основе простых правил, основанных на этих полосах, так и для обучения AI, который, как мы надеемся, впоследствии сможет обобщить этот опыт на большее количество сценариев.

Наш код на данный момент имеет следующий вид:

import numpy as np
from PIL import ImageGrab
import cv2
import time
from directkeys import ReleaseKey, PressKey, W, A, S, D 


def process_img(original_image):
    processed_img = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)
    processed_img = cv2.Canny(processed_img, threshold1=200, threshold2=300)
    return processed_img



def main():
    for i in list(range(4))[::-1]:
        print(i+1)
        time.sleep(1)

    last_time = time.time()
    while(True):
        screen =  np.array(ImageGrab.grab(bbox=(0,40, 800, 640)))
        new_screen = process_img(screen)
        print('Loop took {} seconds'.format(time.time()-last_time))
        last_time = time.time()
        cv2.imshow('window', new_screen)
        #cv2.imshow('window2', cv2.cvtColor(screen, cv2.COLOR_BGR2RGB))
        if cv2.waitKey(25) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break
def process_img(original_image):
    processed_img = cv2.cvtColor(original_image, cv2.COLOR_BGR2GRAY)
    processed_img = cv2.Canny(processed_img, threshold1=200, threshold2=300)
    
    vertices = np.array([[10,500],[10,300],[300,200],[500,200],[800,300],[800,500],
                         ], np.int32)
    processed_img = roi(processed_img, [vertices])
    
    return processed_img 
from IPython.display import Image
Image(filename='idea-of-roi.png')  

На картинке выше приведено схематичное изображение местности, которое мы будем использовать в качестве полигона для нашего автомобиля. Теперь, разумеется, нам нужно написать roi-функцию. ROI (Region of interest — интересующая нас местность).

def roi(img, vertices):
    #blank mask:
    mask = np.zeros_like(img)
    # fill the mask
    cv2.fillPoly(mask, vertices, 255)
    # now only show the area that is the mask
    masked = cv2.bitwise_and(img, mask)
    return masked

def main():
    last_time = time.time()
    while(True):
        screen =  np.array(ImageGrab.grab(bbox=(0,40, 800, 640)))
        new_screen = process_img(screen)
        print('Loop took {} seconds'.format(time.time()-last_time))
        last_time = time.time()
        cv2.imshow('window', new_screen)
        #cv2.imshow('window2', cv2.cvtColor(screen, cv2.COLOR_BGR2RGB))
        if cv2.waitKey(25) & 0xFF == ord('q'):
            cv2.destroyAllWindows()
            break

Image(filename='roi-and-edge.jpg') 

Как вы можете заметить из этого изображения, мы не определили столбы вдоль дороги и провода на них. Также, хотя мы и определили края дороги, дорожная разметка осталась нераспознанной.

Следующая статья — Играем в GTA V c Python. Часть V: алгоритм Хафа.