Spaces:
No application file
No application file
| 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() | |