이제 학습한 모델을 불러들이고 실제 테스트셋을 입력할 경우 Classification으로 정상인지 폐렴인지 학습을 잘해노았는지 검증하여보자.
prediction_model\ = model.load_state_dict(torch.load('./pneumonia_weight.pt'))
def process_image(image_path):
image = Image.open(image_path).convert('RGB') # 이미지를 RGB로 변환
preprocess = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
image_tensor = preprocess(image).unsqueeze(0) # 배치 차원 추가
return image_tensor
- 실제 폐렴으로 진단받은 학습데이를 입력할 경우 폐렴을로 진단이 잘되었다. 뭔가 허전하지만 다음편에는 GradCam으로 어딜보고 폐렴으로 진단했는지 업데이트해볼 예정이다.
# 이미지 경로 설정
image_path = '/home/badagreen/다운로드/archive/chest_xray/test/PNEUMONIA/person1_virus_8.jpeg' # 실제 이미지 경로로 변경하세요.
# 이미지 전처리
img_tensor = process_image(image_path)
# 모델을 통한 예측 수행
model.eval() # 모델을 평가 모드로 설정
with torch.no_grad():
outputs = model(img_tensor.to(device))
_, preds = torch.max(outputs, 1)
predicted_class = preds.item()
# 클래스 레이블 설정 (이 예에서는 0이 '정상', 1이 '폐렴'으로 가정)
classes = ['정상', '폐렴']
print(f"모델 예측: {classes[predicted_class]}")
모델 예측: 폐렴