본문 바로가기
TTL/5주차 ~ 8주차 TIL (9.26 ~ 10.21)

2022-10-13 TIL (yolo를 활용한 사물인식 응용)

by dev_junseok 2022. 10. 13.

https://jundevlog.tistory.com/37

 

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

오늘은 파이토치와 yolo를 이용한 간단한 사물인식 프로그램을 짜보겠습니다. 먼저 pip install -qr https://raw.githubusercontent.com/ultralytics/yolov5/master/requirements.txt # install dependencies 가상..

jundevlog.tistory.com

오늘은 어제 작성한 글을 응용하여 더욱 성능이 향상된 사물인식 모델을 만들어 보도록 하겠습니다.

 

아래의 사진을 다운받아서 맞는 경로에 넣어주도록 합니다.

오늘은 이 사진에서 사람들을 인식한 후,

각 사람들을 잘라서 People1,people2,peopl3... 이런식으로 저장해보도록 하겠습니다.

 

 

 

import torch
import cv2

import ssl
ssl._create_default_https_context = ssl._create_unverified_context

model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
img = cv2.imread('Untitled.jpeg')

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('Untitled.jpeg')
print(tmp_img.shape)

cropped0 = tmp_img[int(result[0][1]):int(result[0][3]), int(result[0][0]):int(result[0][2])]
cropped1 = tmp_img[int(result[1][1]):int(result[1][3]), int(result[1][0]):int(result[1][2])]
cropped2 = tmp_img[int(result[2][1]):int(result[2][3]), int(result[2][0]):int(result[2][2])]
cropped3 = tmp_img[int(result[3][1]):int(result[3][3]), int(result[3][0]):int(result[3][2])]
cropped4 = tmp_img[int(result[4][1]):int(result[4][3]), int(result[4][0]):int(result[4][2])]
print(cropped0.shape)

cv2.imwrite('people01.png', cropped0)
cv2.imwrite('people02.png', cropped1)
cv2.imwrite('people03.png', cropped2)
cv2.imwrite('people04.png', cropped3)
cv2.imwrite('people05.png', cropped4)
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)

각 사람을 cropped0,cropped1,cropped2... 이런식으로 좌표에 맞게 분류해서 저장해줍니다.

그 후 opencv를 이용하여 저장해주면 끝입니다.