인프런 커뮤니티 질문&답변

안녕하세요, 인프런 AI 인턴이에요.

새로운 이미지에서 학습되어 있는 객체를 추출하기 위해서는, 먼저 파이토치(PyTorch) 라이브러리를 이용하여 .pt 파일을 불러온 후, 추론을 수행해야합니다.

다음은, 파이토치를 사용해 .pt 파일을 불러오는 코드입니다.

import torch

# 모델 불러오기
model = torch.load("학습된 모델 경로/모델 이름.pt")

# 모델 추론 모드 설정
model.eval()

.pt 확장자 파일을 불러와 훈련된 모델을 model 변수에 저장합니다. 그 후, model.eval() 메소드를 사용하여 모델을 추론 모드로 설정합니다.

이제 이미지에서 추출하려는 객체에 대한 모델의 출력값이 필요합니다. 이러한 출력값은 모델의 Forward 메소드로 계산됩니다.

다음은 입력 이미지에 대한 추론을 수행하는 코드입니다.

import cv2
import numpy as np

# 이미지 불러오기
image = cv2.imread('이미지 경로/이미지 이름.jpg')

# 이미지 사이즈 조정
resized_image = cv2.resize(image, (640, 640))

# 이미지 배열 초기화
image_tensor = np.zeros((1, 3, 640, 640), dtype=np.float32)
image_tensor[0] = torch.from_numpy(resized_image.transpose(2, 0, 1))

# 모델 입력값 설정
model_input = (image_tensor.cuda(),)

# 모델 추론 실행
output = model(*model_input)

위 코드에서, cv2.imread() 함수를 사용하여 이미지를 읽어옵니다. 그리고, 입력 이미지를 모델에 주입할 수 있도록 사이즈를 조절합니다. YOLOv5 모델은 640x640 크기의 이미지를 입력으로 사용합니다.

이미지는 NumPy 배열로 표현됩니다. 이는 PyTorch의 Tensor로 변환할 수 있는데, 이를 위해 아래 코드와 같이 입력 데이터의 배열을 초기화한 후에, 이미지를 resize() 함수를 이용하여 크기를 맞추고, 이를 0번째 인덱스로 넣어줍니다.

해당 입력값을 모델에 삽입하고 모델의 forward() 메서드를 호출하여 출력값을 계산합니다.

이제 모델을 추론한 결과인 output에서 추출하려는 객체를 찾을 수 있습니다.

출력값을 파싱하는 것은 자세한 모델 구조와 파라미터 등에 따라 다릅니다. YOLOv5에서는 output 변수의 형태를 (N, 6)로 가정합니다. 여기서 N은 이미지에서 추출한 객체의 수입니다. 6개의 값은 [x, y, w, h, cls, conf]의 형태로, 이미지에서 추출한 객체의 좌표와 클래스 정보 등을 나타냅니다.

이제 이러한 값에서 본인이 원하는 객체의 정보를 추출하면 됩니다.

이상입니다. 참고가 되셨길 바랍니다.