[객체탐지] 왕초보도 가능한 인공지능/딥러닝으로 이미지에서 객체 탐지하기(파이썬, pytorch, YOLO v5)

 

 

성능이 정말 좋아서..
이게 안쓸 수가 없잖아?

 

요즘에는 코드 한줄로 딥러닝을 불러올 수 있는 모델들이 많이 소개되고 있습니다.

오늘은 유명한 객체탐지 모델을 사용해보고자 합니다.

환경 설정없이 코드를 테스트해보고 싶다면, 아래 파이토치 공식 홈페이지에서 가능합니다. 

 

https://pytorch.org/hub/ultralytics_yolov5/

불러오는 중입니다...

 

 

 

---

 

 

코드를 좀 살펴볼까요?

YoloV5 를 이용하여 이미지 내의 사람을 탐지하는 코드입니다. 

전체 코드는 20줄이며, 실제로 딥러닝을 실행하는 코드는 아래의 단 3줄입니다.

 

실행방법은 2단계 입니다. 

1. frame = cv2.imread('공원.png') : 이 부분의 이미지를 자신이 원하는 이미지로 변경합니다.
2. 실행, 끝! 

 

자세한 코드의 설명은 주석을 참고 바랍니다. 

 

import cv2
import torch

yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5s',
                            device='cuda:0' if torch.cuda.is_available() else 'cpu')  # 예측 모델
yolo_model.classes = [0]  # 예측 클래스 (사람)
frame = cv2.imread('공원.png') # 입력 이미지 '공원.png' 불러오기
results = yolo_model(frame) # 딥러닝 모델 실행
results_refine = results.pandas().xyxy[0].values # 예측 결과 사용가능하도록 정리 
print(results_refine) # 예측 결과 출력
nms_human = len(results_refine) # 전체 예측 사람 수 계산
print(nms_human) # 전체 사람 수 출력
if nms_human > 0: # 사람 전체 수가 0보다 크면 실행
    for bbox in results_refine: # 예측 결과 하나씩 불러오기
        start_point = (int(bbox[0]), int(bbox[1])) # 객체 영역 오른쪽상단 좌표(x,y)
        end_point = (int(bbox[2]), int(bbox[3])) # 객체 영역 왼쪽하단 좌표(x,y)      
        frame = cv2.rectangle(frame, start_point, end_point, (255, 0, 0), 3) # 이미지에 객체 영역 그리기

while True: # 반복 실행
    cv2.imshow("show image, q를 누르면 종료됩니다", frame) # 이미지 결과 띄우기 
    if cv2.waitKey(1) == ord("q"): #사용자가 키보드로 q를 입력하면 종료
        break

cv2.imwrite("out.png", frame) # 이미지 'out.png' 이름으로 저장하기

 

결과는 다음과 같이 나옵니다.