본문 바로가기
프로젝트 기록/딥러닝 모델 개발_공학설계캡스톤디자인(스마트카ICT)

[YOLOv5, openCV] 딥러닝 모델로 객체detect + openCV bounding box 표시 + 음성 출력

by 소요이 2023. 5. 21.
728x90

모든내용 요약

1. YOLOv5 모델 로드: 이 모델은 훈련된 가중치(best.pt파일)를 사용하며, 이를 이용하여 이미지에서 물체를 감지하는 역할을 함.
2. 경고문구를 설정: 특정 물체가 감지되었을 때 음성으로 재생될 문구 설정
3. 비디오 캡처: 주어진 동영상 파일로부터 프레임을 하나씩 읽어옴 
4. 각 프레임 처리 과정
- 관심 영역 (ROI: Region of Interest) 설정: 물체 감지를 실행할 이미지 영역을 결정. 주어진 점들을 기반으로 다각형을 만들어내고, 이 다각형 내부의 영역을 까맣게 처리함으로써 ROI를 설정하였음
- ROI를 이용 YOLOv5 모델로 물체 감지: 이미지를 RGB 형식으로 변환하여 모델에 전달
- 감지된 물체 바운딩 박스 생성: 감지된 물체의 위치와 크기 정보는 YOLOv5 모델의 결과로부터 가져옴, 이 정보를 이용하여 OpenCV의 rectangle 함수와 putText 함수로 이미지 위에 바운딩 박스와 물체의 클래스 이름을 그림
- 감지된 물체 경고문구 음성 재생: 감지된 물체의 클래스에 따라 해당하는 경고문구를 가져오고, 이를 Google TTS를 이용하여 음성 파일로 변환, pygame 라이브러리를 이용하여 재생
 +) 여러 물체가 동시에 감지된 경우: 스레드를 사용하여 음성을 재생, 각 경고음 사이에 일정 시간을 두어 경고음이 겹치지 않도록 함.
 
요약: 이러한 과정을 반복하여 동영상의 각 프레임에 대해 물체 감지와 경고음 재생을 수행하고, 감지 결과를 시각화함. 이 알고리즘은 YOLOv5, OpenCV, Google TTS, pygame 등 여러 라이브러리를 이용하여 구현하였음
 

사용된 주요 모델

딥러닝 모델 YOLOv5
- 이미지에서 객체를 감지
- model(roi[:, :, ::-1])를 호출하여 관심 영역 (ROI)에서 객체를 감지

감지 결과 시각화 openCV
- bounding box를 그리는데 사용
- OpenCV의 rectangle 함수와 putText 함수를 사용하여 각 객체에 대한 경계 상자를 그리고, 그 위에 객체의 클래스 이름을 출력
 
 

사용한 라이브러리:

모두 아나콘다 터미널에서 pip install로 설치
 
 

파일 디렉토리

 
 

코드 알고리즘

사운드 생성 (google TTS)
- 머리, 손, etc일 때 각 생성되는 파일 명: 'voice_head', 'voice_hands', 'voice_etc'
- 만약 코드가 포함되어 있는 디렉토리에 해당 파일 명을 가진 파일이 있다면 새로운 파일을 만들지 말고 알맞은 파일을 재생만 시키고, 파일이 없다면 gTTS를 이용해 최초 1번만 만들고 그 파일을 재생
 
사운드 재생
- 한 번 객체가 인식되면, 그 때 인식된 객체와 알맞은 음성을 출력
- 인식된 객체와 다른 객체가 인식되면, 새로 인식된 객체에 해당하는 음성을 출력
- 음성 출력이 완료되었는데도 계속 해당 객체가 감지되고 있다면, 1초 기다린 뒤 해당 문구를 출력
- 음성 출력이 되고 있는 와중에, 인식되던 객체가 사라진다면 음성 출력은 0.5초 뒤에 종료
 
 

결과 영상 (동영상)

 

 

코드

"""230521_mp4_say_greenbox.py"""


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

# 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

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

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

    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()

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

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

video.release()
cv2.destroyAllWindows()

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

 


앞으로 보완할 알고리즘 고려사항

 
딥러닝 알고리즘을 사용하여 특정 영역의 객체만 감지하려면, 일반적으로 두 가지 방법이 있습니다:

ROI (관심 영역)에 대한 사전 처리: 이 방법은 OpenCV와 같은 도구를 사용하여 이미지에서 관심 영역을 추출하고, 이 영역에 대해서만 딥러닝 알고리즘을 실행하는 것입니다. 이 방법의 장점은 알고리즘의 실행 속도를 향상시킬 수 있다는 것입니다. 관심 영역 외의 이미지는 처리하지 않기 때문에 계산 효율성이 향상됩니다. 단점은 관심 영역을 잘못 설정하거나, 관심 영역이 이미지와 시간에 따라 변경될 수 있다는 것입니다.

딥러닝 알고리즘으로 ROI를 학습: 이 방법은 관심 영역을 딥러닝 알고리즘에 직접 학습시키는 것입니다. 이 방법의 장점은 훈련 데이터가 충분하다면, 알고리즘이 스스로 관심 영역을 찾아낼 수 있다는 것입니다. 단점은 훈련 데이터를 많이 필요로 하고, 계산 복잡성이 높아질 수 있다는 것입니다.

따라서, 어떤 방법을 선택할지는 프로젝트의 요구 사항에 따라 달라집니다. 시간과 자원이 충분하다면, 두 번째 방법을 사용하여 더욱 강력한 모델을 만들 수 있습니다. 그러나 실시간 처리가 필요하거나, 훈련 데이터가 부족한 경우에는 첫 번째 방법이 더 적합할 수 있습니다.

사전 처리 단계에서 ROI를 설정하는 것이 이중 처리처럼 느껴질 수도 있지만, 딥러닝 모델이 처리해야 하는 데이터의 양을 줄여주고, 따라서 처리 속도를 높여주는 중요한 역할을 합니다. 따라서 어떤 방법을 선택할지는 프로젝트의 목표와 제한 사항에 따라 달라집니다.