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)]
'코딩공부 > Machine Learning' 카테고리의 다른 글
데이터셋 활용하여 리뷰 평가 기능 구현 / TIL_221108 (0) | 2022.11.09 |
---|---|
데이터셋을 활용하여 추천시스템 구현 / TIL_221102 (0) | 2022.11.02 |
머신러닝 라이브러리 / TIL_221012 (0) | 2022.10.12 |
머신러닝 이진 논리회귀 순서 / TIL_221011 (0) | 2022.10.11 |
머신러닝이란? / TIL_221007 (0) | 2022.10.07 |
댓글