Skip to main content

동물 이미지 분류를 손쉽게 구현하는 YOLOv5 완전 가이드 (초보자도 따라 할 수 있는 단계별 튜토리얼)

Published 3 min readMarkdown

# 📌 서론 – “왜 동물 이미지 분류가 필요할까?”

인공지능을 배우는 학생이나 코딩 입문자에게 가장 매력적인 프로젝트 중 하나가 이미지 분류입니다.

특히 동물 사진처럼 직관적이고 재미있는 데이터를 다루면 학습 동기 부여도 확실히 높아지죠.

하지만 “이미지 분류 모델을 만들려면 복잡한 알고리즘을 직접 구현해야 할까?” 라는 고민을 한 번이라도 해보셨나요?

답은 아니오입니다. 최신 객체 탐지 프레임워크인 YOLOv5와 파이토치(PyTorch)만 있으면, 복잡한 전처리와 모델 설계 과정을 크게 단순화할 수 있습니다.

이번 포스트에서는 YOLOv5를 활용해 Animals-10 데이터셋(Kaggle)으로 동물 이미지 분류 파이프라인을 구축하는 전 과정을 초보자도 이해할 수 있도록 차근차근 설명합니다.

👉 핵심 키워드: YOLOv5, 이미지 분류, Animals-10 데이터셋


# 🛠️ 본론

# 1️⃣ 환경 설정 – 파이썬·PyTorch·YOLOv5 설치하기

| 단계 | 명령어 | 설명 |

|------|--------|------|

| 1 | python -V | Python 3.8 이상 확인 |

| 2 | pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu118 | GPU 가속이 가능한 PyTorch 설치 |

| 3 | git clone https://github.com/ultralytics/yolov5.git | YOLOv5 공식 레포 복제 |

| 4 | cd yolov5 && pip install -r requirements.txt | YOLOv5 의존성 설치 |

Tip: CUDA가 지원되는 GPU가 있다면 cu118(CUDA 11.8) 버전을, CPU만 사용할 경우 pip install torch 로 간단히 설치합니다.


# 2️⃣ 데이터 준비 – Animals-10 데이터셋 다운로드 & 구조화

  1. Kaggle 계정을 만든 뒤 Animals-10 데이터셋 페이지에서 Download 버튼을 클릭합니다.

  2. 압축을 풀면 train/, test/ 폴더가 각각 10개의 클래스(예: cat, dog, elephant …) 로 구성됩니다.

  3. YOLOv5는 객체 탐지 형식(이미지 + 라벨 txt) 을 기대하지만, 이미지 분류 용도로는 라벨 파일 없이도 학습이 가능합니다.

    • 이미지 폴더 구조 그대로 data/animals10/ 아래에 train/, val/ 디렉터리를 만들고, 각 클래스 폴더를 복사합니다.
bash

mkdir -p data/animals10/train data/animals10/val

# 80%를 train, 20%를 val 로 분리 (예시)

python split_dataset.py --src data/animals10_original --train data/animals10/train --val data/animals10/val --ratio 0.8

핵심: 라벨 파일이 없더라도 YOLOv5는 폴더 이름을 클래스 이름으로 자동 매핑합니다.


# 3️⃣ YOLOv5 모델 설정 – yolov5s.yaml 수정

YOLOv5는 여러 크기의 사전 학습 모델(s, m, l, x)을 제공합니다. 초보자에게는 가장 가벼운 yolov5s 를 권장합니다.

  1. models/yolov5s.yaml 파일을 열고 nc(class 수) 값을 10 으로 수정합니다.

  2. head 부분은 그대로 두어도 되지만, 이미지 분류에 특화된 cls 헤드만 사용하게 됩니다.

yaml

# models/yolov5s.yaml

nc: 10  # Animals-10 클래스 수

# 4️⃣ 학습 실행 – 한 줄 명령어로 모델 훈련

bash

python train.py \

  --img 640 \

  --batch 32 \

  --epochs 50 \

  --data data/animals10.yaml \

  --cfg models/yolov5s.yaml \

  --weights yolov5s.pt \

  --project runs/train \

  --name animals10_classify
  • --data 옵션에 사용할 yaml 파일(animals10.yaml)은 다음과 같이 작성합니다.
yaml

# data/animals10.yaml

train: ./data/animals10/train

val: ./data/animals10/val

nc: 10

names: ['cat','dog','elephant','giraffe','horse','lion','monkey','panda','rabbit','tiger']

학습이 진행되는 동안 runs/train/animals10_classify 폴더에 loss, precision, recall 그래프와 최종 모델(best.pt)이 저장됩니다.


# 5️⃣ 모델 평가 & 추론 – 테스트 이미지로 결과 확인

bash

python val.py \

  --data data/animals10.yaml \

  --weights runs/train/animals10_classify/weights/best.pt \

  --img 640 \

  --batch 32

평가 결과는 mAP@0.5(mean Average Precision) 로 표시됩니다. Animals-10 데이터셋은 비교적 간단하기 때문에 90% 이상 정확도를 기대할 수 있습니다.

실제 이미지에 적용해 보기

python

import torch

from pathlib import Path

model = torch.hub.load('ultralytics/yolov5', 'custom', path='runs/train/animals10_classify/weights/best.pt')

img = Path('sample_images/elephant.jpg')

result = model(img)

result.print()   # 콘솔에 클래스와 확률 출력

result.show()    # 이미지에 라벨 오버레이 표시

# 📈 결론 – 이제 직접 프로젝트를 확장해 보세요!

  • YOLOv5와 PyTorch만으로 복잡한 이미지 분류 파이프라인을 30분 안에 구축할 수 있었습니다.

  • Animals-10 데이터셋을 활용해 10가지 동물 클래스를 90% 이상 정확도로 구분했으며, 학습·평가·추론 과정을 모두 자동화했습니다.

# 다음 단계 제안

  1. 데이터 증강(albumentations 등)으로 모델 강건성 향상

  2. 다중 클래스(예: 동물+배경) 탐지 모델로 확장

  3. 모바일/Edge 배포(ONNX, TensorRT)로 실시간 애플리케이션 제작

행동 유도: 지금 바로 위 코드를 복사해 로컬 환경에 실행해 보세요! 궁금한 점이나 개선 아이디어가 있다면 댓글로 공유해 주세요. 여러분의 피드백이 다음 포스트를 더욱 풍부하게 만듭니다. 🙌


관련 자료

SEO 키워드 : YOLOv5, 이미지 분류, Animals-10 데이터셋.

No comments yet. Be the first to comment.

More from this blog