[python-arduino] 테스트: 2,3,4보내서 해당번호 LED켜게하는 코드

import cv2
import numpy as np
import torch
import threading
import time
from gtts import gTTS
import os
import pygame
import serial  # serial communication

# YOLOv5 모델 로드
model_path = 'C:/Users/songs/PycharmProjects/mediapipe/yolov5/data/dataset_230515/230518_/best_SY_230518.pt'
model = torch.hub.load('ultralytics/yolov5', 'custom', path=model_path)

# 클래스 이름 불러오기
names = model.module.names if hasattr(model, 'module') else model.names

# 경고문구 설정
warnings = {
    'head': '머리가 감지되었습니다.',
    'hand': "창 밖으로 손을 내밀지 마세요",
    'etc': '모르는 물체가 감지되었습니다.'
}

# 경고음 재생 함수
def play_audio(file):
    pygame.mixer.init()
    pygame.mixer.music.load(file)
    pygame.mixer.music.play()

audio_thread = None
audio_end_time = 0

# Serial port 설정
arduino = serial.Serial('COM6', 9600)  # COM 포트와 baud rate를 아두이노와 동일하게 설정

# 비디오 캡처
video = cv2.VideoCapture('final_test.mp4')
#video = cv2.VideoCapture(1)

try:
    while video.isOpened():
        ret, frame = video.read()
        if not ret:
            break

        #points = np.array([[191, 475], [505, 475], [531, 410], [477, 276], [402, 143], [317, 69], [238, 2], [79, 6], [128, 166], [180, 356]], dtype=np.int32)
        points = np.array([[361, 28], [565, 79], [719, 115], [870, 164], [996, 238],
                           [1069, 350], [1110, 443], [1023, 521], [904, 556], [768, 608],
                           [552, 698], [515, 539], [462, 353], [404, 173]], dtype=np.int32)
        # 관심 영역 설정 (창문 안쪽을 까맣게 처리)
        mask = np.ones_like(frame, dtype=np.uint8) * 255
        cv2.fillPoly(mask, [points], (0, 0, 0))
        roi = cv2.bitwise_and(frame, mask)

        # 객체 탐지 실행
        results = model(roi[:, :, ::-1])  # YOLOv5는 RGB 이미지를 기대하므로 BGR을 RGB로 변환
        labels = results.xyxyn[0][:, -1].cpu().numpy()
        boxes = results.xyxyn[0][:, :-1].cpu().numpy()

        # 바운딩 박스 그리기
        for (x, y, w, h, _), label in zip(boxes, labels):
            x1 = int((x - w / 2) * frame.shape[1])  # frame width
            y1 = int((y - h / 2) * frame.shape[0])  # frame height
            x2 = int((x + w / 2) * frame.shape[1])  # frame width
            y2 = int((y + h / 2) * frame.shape[0])  # frame height
            cv2.rectangle(roi, (x1, y1), (x2, y2), (0, 255, 0), 2)
            cv2.putText(roi, names[int(label)], (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 2)

        # 새로운 객체 탐지 시 음성 재생 및 아두이노로 시리얼 통신
        if len(labels) > 0 and time.time() > audio_end_time:
            detected_classes = [names[int(label)] for label in labels]
            for detected_class in detected_classes:
                warning = warnings.get(detected_class, warnings['etc'])
                audio_file = f'voice_{detected_class}.mp3'
                if not os.path.isfile(audio_file):  # 파일이 존재하지 않는 경우만 TTS를 이용하여 생성
                    # TTS를 이용하여 경고문장을 음성으로 변환
                    tts = gTTS(warning, lang='ko')
                    tts.save(audio_file)

                # 음성 재생 시작 (새로운 스레드에서)
                if audio_thread is not None:
                    audio_thread.join()
                audio_thread = threading.Thread(target=play_audio, args=(audio_file,))
                audio_thread.start()

                # 아두이노로 객체 이름 시리얼 통신
                if detected_class == 'head':
                    arduino.write('2'.encode())
                elif detected_class == 'hand':
                    arduino.write('3'.encode())
                else:
                    arduino.write('4'.encode())

                # 경고음이 끝날 때까지 기다리는 시간 설정
                audio_end_time = time.time() + 3.5  # 가정: 각 경고음은 대략 3.5초 걸림

        # 아두이노로 아무것도 감지되지 않았음을 시리얼 통신
        if len(labels) == 0:
            arduino.write('5'.encode())

        cv2.imshow('ROI', roi)
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

finally:  # 프로그램이 종료되기 전에 시리얼 포트를 닫는다.
    video.release()
    cv2.destroyAllWindows()

    # 마지막으로 재생된 오디오가 끝날 때까지 기다림
    if audio_thread is not None:
        audio_thread.join()

    # 시리얼 포트 닫기
    arduino.close()