import cv2
import numpy as np
import math

# Função para calcular o ângulo entre duas linhas
def calcular_angulo(linha1, linha2):
    angulo1 = math.atan2(linha1[1] - linha1[3], linha1[0] - linha1[2])
    angulo2 = math.atan2(linha2[1] - linha2[3], linha2[0] - linha2[2])
    angulo = math.degrees(angulo1 - angulo2)
    angulo = (angulo + 360) % 360
    if angulo > 180:
        angulo -= 360
    return angulo

# Função para isolar a cor amarela em uma imagem
def isolar_cor_amarela(frame):
    # Convertendo para o espaço de cores HSV
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)

    # Definindo o intervalo da cor amarela em HSV
    # Estes valores podem precisar ser ajustados dependendo do tom da cor amarela
    lower_yellow = np.array([20, 100, 100])
    upper_yellow = np.array([30, 255, 255])

    # Criando uma máscara que isola a cor amarela
    mask = cv2.inRange(hsv, lower_yellow, upper_yellow)

    return mask

# Função para processar cada frame do vídeo
def processar_frame(frame):
    altura, largura = frame.shape[:2]
    centro_imagem_x = largura // 2

    # Isolar a cor amarela
    mask_amarela = isolar_cor_amarela(frame)

    # Encontrar os contornos na máscara de cor amarela
    contornos, _ = cv2.findContours(mask_amarela, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    # Desenhar uma linha de orientação no centro da imagem
    cv2.line(frame, (centro_imagem_x, 0), (centro_imagem_x, altura), (255, 0, 0), 2)

    # Supondo que as duas maiores áreas amarelas são as bordas da rua
    if len(contornos) >= 2:
        # Encontrar os dois maiores contornos que serão as bordas
        bordas = sorted(contornos, key=cv2.contourArea, reverse=True)[:2]

        # Calcular a linha média entre as duas bordas amarelas
        linhas_bordas = []
        for borda in bordas:
            vx, vy, x0, y0 = cv2.fitLine(borda, cv2.DIST_L2, 0, 0.01, 0.01)
            m = vy / vx
            b = y0 - m * x0
            linhas_bordas.append((m, b[0]))

        # Calcular a linha central como média das linhas das bordas
        m_centro = np.mean([line[0] for line in linhas_bordas])
        b_centro = np.mean([line[1] for line in linhas_bordas])
        y1_centro = int(m_centro * 0 + b_centro)
        y2_centro = int(m_centro * largura + b_centro)

        # Garantir que os pontos são inteiros e dentro dos limites da imagem
        y1_centro = max(0, min(y1_centro, altura))
        y2_centro = max(0, min(y2_centro, altura))

        cv2.line(frame, (0, y1_centro), (largura, y2_centro), (0, 255, 255), 2)

        # Calcular o ângulo entre a linha central de orientação e a linha média das bordas
        linha_central = [centro_imagem_x, 0, centro_imagem_x, altura]
        linha_rua_centro = [0, y1_centro, largura, y2_centro]
        angulo = calcular_angulo(linha_central, linha_rua_centro)

        return frame, angulo

    return frame, None

# Inicializando a captura de vídeo
cap = cv2.VideoCapture(1)  # Substitua 1 pelo índice correto da sua câmera

try:
    while True:
        # Capturando frame a frame
        ret, frame = cap.read()
        if not ret:
            print("Falha ao capturar imagem da câmera")
            break

        # Processar o frame atual
        frame_processado, angulo = processar_frame(frame)

        if angulo is not None:
            print(f"Ângulo do caminho com relação à linha central: {angulo:.2f} graus")

        # Exibir o frame processado
        cv2.imshow('Frame Processado', frame_processado)

        # Pressione 'q' para sair do loop
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
finally:
    # Quando tudo estiver feito, liberar a captura
    cap.release()
    cv2.destroyAllWindows()