app.py

import pygame
import cv2
import time
import serial
import sys

from ultrafast.inference_onnx import AI_TRT
from classification.inference_onnx import inference_classification
from setting_AI import *
from a_utils_func_2_model import (
    CLEAN_DATA_CSV_DIRECTION,
    CLEAN_DATA_CSV_DIRECTION_STRAIGHT,
)
from a_control_classification import USE_CLASSIFICATION



serial_p = False
if serial_p:
    serial_port = serial.Serial(
        "COM8", 9600, serial.EIGHTBITS, serial.PARITY_NONE, serial.STOPBITS_ONE
    )

# Initialize camera
cap = cv2.VideoCapture("./videos/test_video.mp4")
cap_ = cv2.VideoCapture("./videos/test_video.mp4")
if serial_p:
    cap = cv2.VideoCapture(0, cv2.CAP_DSHOW)
    cap_ = cv2.VideoCapture(1, cv2.CAP_DSHOW)

pygame.init()

# Screen settings
screen_width, screen_height = 1600, 900
screen = pygame.display.set_mode((screen_width, screen_height))
pygame.display.set_caption("AI Camera Control")

# Colors
WHITE = (240, 240, 240)
GREEN = (34, 177, 76)
RED = (200, 50, 50)
BLACK = (20, 20, 20)
GRAY = (180, 180, 180)
DARK_GRAY = (100, 100, 100)

# Fonts
font = pygame.font.Font(None, 50)
small_font = pygame.font.Font(None, 36)

# Buttons
start_button = pygame.Rect(100, 820, 220, 70)
end_button = pygame.Rect(400, 820, 220, 70)

# Slider settings
ROTATION_SPEED = 10
slider_rect = pygame.Rect(750, 850, 300, 10)
slider_knob_rect = pygame.Rect(750 + int((ROTATION_SPEED / 30) * 300) - 10, 840, 20, 30)
slider_dragging = False

DISPLAY_WIDTH = 600  # Fixed width for display
DISPLAY_HEIGHT = 400
time.sleep(1)

time_stop = sys.maxsize
sleep_time = sys.maxsize
running = True
active = False
clear = True
push_results = []

while running:
    start_time = time.time()
    screen.fill(WHITE)
    pygame.draw.rect(screen, DARK_GRAY, (0, 800, screen_width, 100))

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        elif event.type == pygame.MOUSEBUTTONDOWN:
            if start_button.collidepoint(event.pos):
                active = True
                clear = True
            elif end_button.collidepoint(event.pos):
                active = False
                print("Stopped pushing")
            elif slider_knob_rect.collidepoint(event.pos):
                slider_dragging = True
        elif event.type == pygame.MOUSEBUTTONUP:
            slider_dragging = False
        elif event.type == pygame.MOUSEMOTION and slider_dragging:
            slider_knob_rect.x = max(
                slider_rect.x, min(event.pos[0] - 10, slider_rect.x + 300 - 20)
            )
            ROTATION_SPEED = int(((slider_knob_rect.x - slider_rect.x) / 300) * 50)

    _, frame = cap.read()
    _, frame_ = cap_.read()
    frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    visualization_img = frame

    inference_classification(frame_)

    if active:
        if USE_CLASSIFICATION.check() == "STRAIGHT":
            if clear:
                CLEAN_DATA_CSV_DIRECTION()
                CLEAN_DATA_CSV_DIRECTION_STRAIGHT()
                clear = False
            visualization_img, PUSH_RETURN, Have_lane = AI_TRT(
                frame, paint=True, resize_img=True
            )
            if PUSH_RETURN:
                if serial_p:
                    serial_port.write(PUSH_RETURN.encode())
                push_results.append(PUSH_RETURN)
                if len(push_results) > 5:
                    push_results.pop(0)
                angle = min(30, int(PUSH_RETURN.split(":")[1]))
                sleep_time = angle / ROTATION_SPEED
                time_stop = time.time()

            if time.time() - time_stop >= sleep_time:
                if serial_p:
                    serial_port.write(PUSH_STOP.encode())
                push_results.append(PUSH_STOP)
                if len(push_results) > 5:
                    push_results.pop(0)
                time_stop = sys.maxsize

        else:

            hard_left = "X:000"
            hard_right = "Y:000"

            hard_time_1 = 0.5
            hard_time_2 = 1.5
            hard_time_3 = 0.5

            if USE_CLASSIFICATION.check() == "LEFT" and serial_p:
                serial_port.write(hard_right.encode())
                print(hard_right)
                time.sleep(0.5)
                serial_port.write(PUSH_STOP.encode())
                print(PUSH_STOP)

                serial_port.write(hard_left.encode())
                print(hard_left)
                time.sleep(1.5)
                serial_port.write(PUSH_STOP.encode())
                print(PUSH_STOP)

                serial_port.write(hard_right.encode())
                print(hard_right)
                time.sleep(0.5)
                serial_port.write(PUSH_STOP.encode())
                print(PUSH_STOP)

                USE_CLASSIFICATION.change("STRAIGHT")

            if USE_CLASSIFICATION.check() == "LEFT" and serial_p:
                serial_port.write(hard_left.encode())
                print(hard_left)
                time.sleep(0.5)
                serial_port.write(PUSH_STOP.encode())
                print(PUSH_STOP)

                serial_port.write(hard_right.encode())
                print(hard_right)
                time.sleep(1.5)
                serial_port.write(PUSH_STOP.encode())
                print(PUSH_STOP)

                serial_port.write(hard_left.encode())
                print(hard_left)
                time.sleep(0.5)
                serial_port.write(PUSH_STOP.encode())
                print(PUSH_STOP)

                USE_CLASSIFICATION.change("STRAIGHT")

    text_cls = font.render(USE_CLASSIFICATION.check(), True, (0, 0, 255))
    text_rect = text_cls.get_rect(center=(900, 200))
    screen.blit(text_cls, text_rect)

    elapsed_time = time.time() - start_time
    fps = 1 / elapsed_time if elapsed_time > 0 else 0

    text_fps = font.render(f"FPS: {fps:.2f}", True, (0, 0, 255))
    text_rect = text_fps.get_rect(center=(900, 250))
    screen.blit(text_fps, text_rect)

    visualization_img = cv2.resize(
        visualization_img,
        # (visualization_img.shape[1] // 2, visualization_img.shape[0] // 2),
        (DISPLAY_WIDTH, DISPLAY_HEIGHT),
    )

    pygame_frame = pygame.surfarray.make_surface(
        cv2.rotate(cv2.flip(visualization_img, 1), cv2.ROTATE_90_COUNTERCLOCKWISE)
    )
    screen.blit(pygame_frame, (10, 10))

    # Buttons
    pygame.draw.rect(screen, GREEN if active else GRAY, start_button, border_radius=15)
    pygame.draw.rect(screen, RED, end_button, border_radius=15)
    screen.blit(
        font.render("Start", True, WHITE), (start_button.x + 70, start_button.y + 20)
    )
    screen.blit(font.render("End", True, WHITE), (end_button.x + 80, end_button.y + 20))

    # Slider
    pygame.draw.rect(screen, GRAY, slider_rect, border_radius=5)
    pygame.draw.ellipse(screen, BLACK, slider_knob_rect)
    screen.blit(font.render(f"Speed: {ROTATION_SPEED}", True, WHITE), (1080, 820))

    # Display push results
    for i, result in enumerate(reversed(push_results)):
        screen.blit(small_font.render(result, True, BLACK), (1200, 600 - i * 40))

    pygame.display.flip()

cap.release()
cv2.destroyAllWindows()
pygame.quit()