머신러닝 yolo를 이용한 이미지에서 사람 인식 / TIL_221013

import torch
import cv2

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

img = cv2.imread('zidane.jpg')
# 폴더 안에 있는 이미지 파일을 읽어 온다.
results = model(img)
# 위에서 불러온 이미지를 모델에 담아서
results.save()
# 저장
print(results)
# 1

result = results.pandas().xyxy[0].to_numpy()
# 넘파이 배열로 바꿔 result에 담아주고
print(result)
# 2
result = [item for item in result if item[6]=='person']
# result의 person만 저장
print(result)
# 3

tmp_img = cv2.imread('zidane.jpg')
cropped = tmp_img[int(result[0][1]):int(result[0][3]), int(result[0][0]):int(result[0][2])]
# 원본 이미지를 다시 불러 result에 저장 된 xmin, ymin, xmax, ymax 값으로 자름
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())), (0,0,255))
# 원본 이미지를 다시 불러 result에 저장 된 xmin, ymin, xmax, ymax 값으로 박스를 그림
cv2.imwrite('result.png', tmp_img)

 

# 1 출력값

model에 이미지를 담고 저장하니 바로 이미지를 분석하여 2명의 사람과 1개의 타이를 인식

image 1/1: 720x1280 2 persons, 1 tie
Speed: 2.0ms pre-process, 112.7ms inference, 2.0ms NMS per image at shape (1, 3, 384, 640)

# 2 출력값

[[743.7371826171875 45.197723388671875 1149.8685302734375 720.0 0.8572446703910828 0 'person']
 [120.3214111328125 197.51776123046875 853.21484375 718.7362060546875 0.6205654740333557 0 'person']
 [439.6195983886719 437.1634521484375 500.7126770019531 711.709716796875 0.2528429627418518 27 'tie']]

# 3 출력값

[array([743.7371826171875, 45.197723388671875, 1149.8685302734375, 720.0, 0.8572446703910828, 0, 'person'], dtype=object),
array([120.3214111328125, 197.51776123046875, 853.21484375, 718.7362060546875, 0.6205654740333557, 0, 'person'], dtype=object)]