Python (29) - YOLOv5 모델을 이용한 애니메이션 이미지를 검출해보자

YOLO 모델을 이용하여 애니메이션 이미지를 찾는 재미있는 프로젝트를 발견하여 업로드합니다.

 

YOLO (You Only Live Once)

탐지된 이미지 - darknet

YOLO(You Only Look Once)는 실시간 객체 검출 시스템입니다. 이름과 마찬가지로 이미지를 한 번 보는 것만으로 객체를 탐지하고 분류할 수 있습니다. YOLO는 빠르고 간단하고 정확하여 많이 사용 중인 모델입니다. 드론, 자율주행, 보안 카메라 등 실시간으로 객체 탐지가 필요한 모든 분야에서 사용 중입니다. 

주요 특징은 다음과 같습니다.

• 실시간으로 동작 : 기존의 객체 검출 모델보다 훨씬 빠르며 실시간 작업에 적합합니다.
• 원본 이미지 그대로 분석 : 이미지를 나누고 분석하는 방식이 아닌 한 번에 통과시켜 객체를 탐지하는 방식입니다. 
• 객체의 위치와 종류 동시 탐지 : 단일 신경망에서 객체의 위치와 종류를 동시에 탐지하여 간단하고 효율적입니다.

현재는 YOLO 11까지 나왔으며 아래는 홈페이지 링크입니다.

https://docs.ultralytics.com/ko#overview

홈

 

애니메이션 이미지 검출하기

애니메이션 이미지

해당 프로젝트는 yolov5 모델에다가 픽시브에 있는 데이터로 학습을 진행하였습니다. 해당 프로젝트를 이용하려고 합니다.

Pytorch 호환 이슈가 있기 때문에 venv등 가상 환경을 만드는 것을 권장합니다.

https://github.com/zymk9/yolov5_anime

GitHub - zymk9/yolov5_anime: Yet another anime face detector based on yolov5.

 

설치 및 테스트

1. Repo 복제 및 패키지 설치

Github에서 저장소를 복제하고 필요한 의존성 패키지를 설치합니다.

git clone https://github.com/zymk9/yolov5_anime.git 
cd yolov5_anime
pip install -r requirements.txt

 

2. 가중치 파일 다운로드

아래 코드를 실행시켜 Google Drive에서 다운로드 합니다.

설치된 프로젝트 안에 utils을 이용하기 때문에 Repo가 있는 디렉터리를 참조해서 다운로드 합니다.

from utils.google_utils import gdrive_download 
gdrive_download('1-MO9RYPZxnBfpNiGY6GdsqCeQWYNxBdl','yolov5x_anime.pt')

직접 설치

 

3. 탐지 테스트 실행 (detect.py 실행)

import subprocess

subprocess.run([
    'python',
    './yolov5_anime/detect.py',
    '--weights', './yolov5_anime/yolov5x_anime.pt',
    '--source', './miku.jpg',
    '--output', './result'
])

여기서 argument는 다음과 같습니다.

• --weights: 사용할 모델 가중치 파일 경로
• --source: 입력 이미지 파일
• --output: 결과를 저장할 출력 폴더 경로

 

4. 결과 확인

참고: 미쿠는 애니메이션이 아닙니다. 보컬로이드 입니다.

예시 결과 - 하츠네 미쿠

실행 후 ./result 폴더를 확인하면 애니메이션 캐릭터 얼굴의 탐지 여부를 확인 할 수 있습니다. 확인 이미지는 다음과 같이 나옵니다.

 

5. 해당 결과를 활용하고 싶다면?

import subprocess
import glob

# YOLOv5 탐지
subprocess.run([
    'python',
    './yolov5_anime/detect.py',
    '--weights', './yolov5_anime/yolov5x_anime.pt',
    '--source', './miku.jpg',
    '--output', './result',
    '--save-txt'
])

# 결과 폴더 내 txt 파일 여부
txt_files = glob.glob('./result/**/*.txt', recursive=True)

if txt_files:
    print("Miku appear")
else:
    print("Miku disappear")

detect.py를 직접 수정 하는 방법 등 여러 가지가 있으나 yolov5 argument중에 save-txt가 있습니다. 검출되면 txt 파일이 나오며 없는 경우에는 나오지 않습니다. 이를 이용한 예시입니다.