python에서 yolov5 모델을 이용해보기

읽을거리&정보공유
, , ,
1

python에서 yolov5 모델을 이용해 보는 방법을 공유합니다
모델은 pytorch hub에서 import해서 사용해보겠습니다.
PC에서 수행해 본 결과지만, GPU를 사용하지 않아서 모바일 및 다른데서도 동일하게 적용 가능한 방법입니다.

먼저 torch 및 관련 모듈을 pip로 설치합니다

pip install ipython seaborn matplotlib torchvision tqdm pillow pyyaml requests pandas opencv-python torch protobuf psutil

추가로 주피터 노트북이 필요하다면 설치합니다

pip install jupyter

그리고 아래처럼 허브에서 필요한 모듈을 import해오고 실행을 해봅니다
2shot.png는 python에서 youtube 화면 screenshot 저장하기 게시물을 참조하면 저장할수 있습니다.
이외에 다른 url을 붙여서 활용도 가능합니다

#!/usr/bin/env python3

import torch

# Model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s') 

# Images
#img = 'https://ultralytics.com/images/zidane.jpg'  # or file, Path, PIL, OpenCV, numpy, list
img = './2shot.png'

# Inference
results = model(img)

# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.
results.save()

결과.print()를 하면, 아래와 같은 결과로 알려주며
결과.save()를 하면 ./runs/detect/exp{실험횟수} 의 디렉토리에 결과를 저장합니다

Using cache found in /home/user/.cache/torch/hub/ultralytics_yolov5_master
YOLOv5 🚀 2022-7-2 Python-3.10.5 torch-1.12.0+cu102 CPU

Fusing layers... 
YOLOv5s summary: 213 layers, 7225885 parameters, 0 gradients
Adding AutoShape... 
image 1/1: 454x806 2 persons, 3 cars, 1 bus
Speed: 21.7ms pre-process, 130.6ms inference, 1.2ms NMS per image at shape (1, 3, 384, 640)
Saved 1 image to runs/detect/exp2

save된 이미지는 아래와 같습니다
2shot

작은 모델(yolov5s)이라 확실한 부분만 판단해둔것을 볼 수 있습니다
yolov5 가 지원하는 더 큰모델(yolov5x6)로 바꿔보겠습니다
또한 성능 향상을 위해 inference시에 추가로 data augment도 하도록 추가해보겠습니다.

model과 inference하는 부분을 바꿔줍니다

...
model = torch.hub.load('ultralytics/yolov5', 'yolov5x6') 
...
results = model(img, augment=True)
...

실행하면 결과가 아래와 같이 나옵니다

Using cache found in /home/user/.cache/torch/hub/ultralytics_yolov5_master
YOLOv5 🚀 2022-7-2 Python-3.10.5 torch-1.12.0+cu102 CPU

Downloading https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5x6.pt to yolov5x6.pt...
100%|██████████████████████████████████████████████████████████████████████| 270M/270M [01:22<00:00, 3.41MB/s]

Fusing layers... 
YOLOv5x6 summary: 574 layers, 140730220 parameters, 0 gradients
Adding AutoShape... 
image 1/1: 454x806 10 persons, 4 cars, 1 motorcycle, 2 buss, 1 truck
Speed: 21.4ms pre-process, 2267.2ms inference, 2.5ms NMS per image at shape (1, 3, 384, 640)
Saved 1 image to runs/detect/exp3

yolov5의 여러 버전 중에서 가장 크고 성능 좋은 모델을 사용했으므로,
아까보다 더 많은 사물을 탐지한것을 볼수 있습니다
결과 사진은 아래와 같습니다
2shot

yolov5도 pytorch hub에서 모델을 import해서 편리하게 사용가능합니다.
해당 모델의 보다 자세한 내용은 GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite 참조하시면 됩니다

감사합니다

3개의 좋아요