2022-10-12 TIL (사물 인식)

오늘은 파이토치와 yolo를 이용한 간단한 사물인식 프로그램을 짜보겠습니다.

 

먼저

pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt # install dependencies

가상환경에 접속 후, 해당 코드를 입력하여 패키지들을 설치해 줍니다.

 

그 다음, https://ultralytics.com/images/zidane.jpg 해당 링크로 접속해서 사진을 pc에 저장합니다.(zidand.jpg)입니다.

 

 

import torch
import cv2

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
imgs = ['https://ultralytics.com/images/zidane.jpg']  # batch of images

results = model(imgs)

results.save() #or. show()

 

해당 코드를 입력 후, 실행해줍니다.

 

실행이 되셨으면, 작업 디렉터리에 runs 파일이 생성되 있는것을 볼 수 있습니다.

 

안에 경로로 들어가보시면

해당 사진과 같이 사물을 인식하는것을 볼 수 있습니다.

 

import torch
import cv2
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

img = cv2.imread('zidane.jpg')
results = model(img)
results.save()

result = results.pandas().xyxy[0].to_numpy()
result = [item for item in result if item[6]=='person']

tmp_img = cv2.imread('zidane.jpg')
print(tmp_img.shape)
cropped = tmp_img[int(result[0][1]):int(result[0][3]), int(result[0][0]):int(result[0][2])]
print(cropped.shape)
cv2.imwrite('result2.png', cropped)
cv2.rectangle(tmp_img, (int(results.xyxy[0][0][0].item()), int(results.xyxy[0][0][1].item())), (int(results.xyxy[0][0][2].item()), int(results.xyxy[0][0][3].item())), (255,255,255))
cv2.imwrite('result.png', tmp_img)

해당 코드를 입력하면 result.png, result2.png 파일이 생성 되는 것을 볼 수 있습니다.