ビデオ認識で走行

import cv2
import numpy as np

from carCtroller import CarController
from carCtroller import MOTOR_FORWARD, MOTOR_REVERSE

import time
import dip

THETA_MIN = 0.6
THETA_MAX = 2.4

ctr = CarController()

capture = cv2.VideoCapture(-1)
capture.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

fourcc = cv2.VideoWriter_fourcc(*'DIVX')
vdwriter =cv2.VideoWriter('./videos/cosmos_load.avi', fourcc, 18, (640, 480))
vdwriter2 = cv2.VideoWriter('./videos/cosmos_load_line.avi', fourcc, 18, (640, 480))

font = cv2.FONT_HERSHEY_SIMPLEX

speed = 155

frame_cnt = 0

ctr.set_wheel_direction(85)
time.sleep(1)
ctr.play_buzzer()
time.sleep(1)
ctr.set_camera_pan(89)
ctr.play_buzzer()
time.sleep(1)
ctr.play_buzzer()



try:
    while True:
        #key = cv2.waitKey(10)
        # print(key)
        
        ret, frame = capture.read()
        if ret:
            frame = cv2.flip(frame, -1) 
            '''
            res = dip.detect_red_light(frame)
            if res:
                ctr.set_motor_stop(MOTOR_FORWARD, 0)
                print("Red light - stop")
                
                # 5. 대기하면서 원이 없어지는지 검사한다.
                text = "Red light - stop"
                cv2.putText(frame, text, (10,450), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
    
                vdwriter.write(frame)
                vdwriter2.write(frame)
                continue
            ctr.set_motor_velocity(MOTOR_FORWARD, speed)
            '''
            img_gray = dip.preprocessing(frame)        
            img_canny = dip.canny(img_gray)
            height, width = frame.shape[:2]

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

            print(frame_cnt, end=', ')
            result = frame.copy()
            text = str(frame_cnt)
            #cv2.putText(result, str(ctr.wheel_dirction)+' '+ str(frame_cnt), (10,450), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
            rho_theta_left = dip.find_rho_theta_left(img_left)
            if rho_theta_left is not None:
                result = dip.draw_line(result, rho_theta_left[0], rho_theta_left[1])
                #print(rho_theta_left[1], end=', ')
                text += ' ' + str(rho_theta_left[1])
            else:
                #result = frame
                #print('None', end=', ')
                text += ' None'
            
            
            rho_theta_right = dip.find_rho_theta_right(img_right)
            if rho_theta_right is not None:
                result = dip.draw_line(result, rho_theta_right[0], rho_theta_right[1])
                #print(rho_theta_right[1], end=', ')
                text += ' ' + str(rho_theta_right[1])
            else:
                #print('None', end=', ')
                text += ' None'

            print()
            '''
            if key == 119:      # w
                ctr.set_motor_velocity(MOTOR_FORWARD, speed)
            elif key == 115:    # s
                ctr.set_motor_stop(MOTOR_FORWARD, 0)
            elif key == 120:    # x
                ctr.set_motor_velocity(MOTOR_REVERSE, speed)
            elif key == 100:    # d
                ctr.set_wheel_direction(50)
            elif key == 97:     # a
                ctr.set_wheel_direction(130)
            elif key == 98:     # b
                ctr.play_buzzer()
            elif key == 105:    # i
                ctr.set_camera_tilt(-10)
            elif key == 44:     # <
                ctr.set_camera_tilt(10)
            elif key == 106:    # j
                ctr.set_camera_pan(10)
            elif key == 108:    # l
                ctr.set_camera_pan(-10)
            elif key == 101:
                speed += 5
            elif key == 113:
                speed -= 5
            elif key == -1:
                #ctr.set_motor_velocity(MOTOR_FORWARD, 0)
                pass
            cv2.imshow('test', result)
            '''
            if frame_cnt < 3:
                #ctr.set_motor_velocity(MOTOR_FORWARD, speed)
                pass
            
            if rho_theta_left is None and rho_theta_right is None:
                #ctr.set_motor_stop(MOTOR_FORWARD, 0)
                pass
                
            elif rho_theta_left is None:
                if rho_theta_right[1] > THETA_MIN and rho_theta_right[1] < THETA_MAX:
                    ctr.set_wheel_direction(70)
                elif rho_theta_right[1] >= THETA_MAX:
                    ctr.set_wheel_direction(60)
                elif rho_theta_right[1] <= THETA_MIN:
                    ctr.set_wheel_direction(120)
                    
            elif rho_theta_right is None:
                if rho_theta_left[1] > THETA_MIN and rho_theta_left[1] < THETA_MAX:
                    ctr.set_wheel_direction(110)
                elif rho_theta_left[1] >= THETA_MAX:
                    ctr.set_wheel_direction(60)
                elif rho_theta_left[1] <= THETA_MIN:
                    ctr.set_wheel_direction(120)
            else:
                
                if rho_theta_left[1] >= THETA_MAX or rho_theta_right[1] >= THETA_MAX:
                    ctr.set_wheel_direction(60)
                elif rho_theta_left[1] <= THETA_MIN or rho_theta_right[1] <= THETA_MIN:
                    ctr.set_wheel_direction(120)
                else:
                    ctr.set_wheel_direction(90)
            
            frame_cnt += 1
            text +=  ' ' + str(ctr.wheel_direction)
            cv2.putText(result, text, (10,450), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
    
            vdwriter.write(frame)
            vdwriter2.write(result)
            #cv2.imshow('result', result)
        #if key == 27:
        #   time.sleep(1)
        #  ctr.set_motor_stop(MOTOR_REVERSE, 0)
        # time.sleep(1)
        #break
    # end-of-while
except KeyboardInterrupt:
    ctr.set_motor_stop(MOTOR_FORWARD, 0)
ctr.set_motor_stop(MOTOR_FORWARD, 0)
capture.release()
vdwriter.release()
cv2.destroyAllWindows()