[딥러닝] Fast API를 활용한 유방 내 종양 세그멘테이션 모델 배포 (1)

김영민·2024년 9월 3일

DeepLearning

목록 보기

33/33

석사 과정 중 딥러닝 모델을 개발하고 연구하는 것은 많이 진행하였으나, 모델을 배포하는 경험이 없어서 이번 기회에 진행해보고자 한다.
사용할 모델은 석사 과정에서 개발한 유방 MRI 내 종양을 세그멘테이션하고 종양이 존재할 확률을 예측하는 모델이다. (classification & segmentation)

해당 프로젝트는 챗지피티와 함께 진행하였다.
~~챗지피티와 함께라면...모든지 가능하니까...~~

FastAPI ?

FastAPI는 Python 기반의 고성능 웹 프레임워크로, 딥러닝 모델을 RESTful API로 손쉽게 배포할 수 있게 도와준다.

배포 순서

환경 설정
딥러닝 모델 준비
FastAPI 애플리케이션 작성
API 테스트

1. 환경 설정

1.1 필요한 패키지 설치

Fast API를 사용하기 위해 필요한 패키지를 설치한다.

pip install fastapi uvicorn pydicom

나는 추가적으로 dicom 이미지를 전처리해야 하므로 pydicom도 설치하였다.

1.2 프로젝트 구조

모델 아키텍처에 사용되는 파이썬 파일은 CBAM.py, gmic_v2.py, modules_v2.py, tools.py 이다.
main.py에서 FastAPI 애플리케이션을 작성한다.

2. 딥러닝 모델 준비

딥러닝 모델 구조와 이미 학습된 가중치를 저장한 Breast_MTL.pth 파일을 통해 모델을 불러올 수 있다.

3. FastAPI 애플리케이션 작성

3.1 라이브러리 임포트 및 FastAPI 애플리케이션 생성

from fastapi import FastAPI, HTTPException, File, UploadFile
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel
from PIL import Image
import torch
import os
import cv2
from .model import load_model
from .data_utils import *


app = FastAPI(title="딥러닝 모델 API")

3.2 모델 불러오기


# 모델 로드
model = load_model("/Users/kimyoungmin/my_fastapi_app/models/Breast_MTL.pth")

모델을 로드합니다.

3.3 Dicom 이미지 처리 및 예측 수행

@app.post("/process/")
async def process_dicom(file: UploadFile = File(...)):
    try:
        # DICOM 파일을 서버에 저장
        dicom_file_path = f"/Users/kimyoungmin/my_fastapi_app/input_data/{file.filename}"
        
        with open(dicom_file_path, "wb") as f:
        	f.write(await file.read())

        # DICOM 이미지 로드 및 전처리
        image_array = load_dicom_image(dicom_file_path)

        dcm_forOverlay = (image_array.squeeze()*255.0).astype(np.uint8)
        dcm_forOverlay = np.expand_dims(dcm_forOverlay, axis=2)
        dcm_forOverlay = cv2.cvtColor(dcm_forOverlay, cv2.COLOR_GRAY2BGR)

        # 모델에 입력할 형태로 변환 (예: 배치 크기 추가)
        input_tensor = torch.tensor(image_array, dtype=torch.float32)

        # 모델 추론
        with torch.no_grad():
            _, _, prediction, output_img = model(input_tensor)
        prediction = prediction.item()

        output_img = torch.sigmoid(output_img).squeeze().cpu().numpy()
        output_img = (output_img * 255.0).astype(np.uint8)
        output_img = np.expand_dims(output_img, axis=2)
        output_img = cv2.applyColorMap(output_img, cv2.COLORMAP_JET)
        output_img[:, :, 0] = 0  # 빨간색으로 설정
        output_img = output_img[:, :, ::-1]

        alpha = 0.7  # 원본 이미지와 overlay 이미지의 가중치 조절
        overlay_result = cv2.addWeighted(dcm_forOverlay, alpha, output_img, 1 - alpha, 0)

        # 결과 PNG 파일 저장
        png_file_path = f"/Users/kimyoungmin/my_fastapi_app/output_data/{os.path.splitext(file.filename)[0]}.png"
        save_png(overlay_result, png_file_path)

        prediction = f"{prediction:.4f}"
        # 실수 값과 PNG 파일 경로 반환
        return JSONResponse({
            "prediction": prediction,
            "png_file": png_file_path
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

post를 통해 dicom 파일 불러오기.
출력을 위한 output_img 구현 진행.
- 기존 이미지에 종양 예측 세그멘테이션 부위를 overlay 하도록 구현.
이 과정에서 오류가 발생하면 500 에러 출력
최종 결과물은 output_data라는 디렉토리에 저장될 것이다.

3.4 PNG 파일 다운로드

@app.get("/download_png/")
async def download_png(file_path: str):
    if os.path.exists(file_path):
        return FileResponse(file_path, media_type="image/png", filename=os.path.basename(file_path))
    else:
        raise HTTPException(status_code=404, detail="File not found")

다운로드 받을 수 있도록 진행

4. API 테스트

터미널에서 프로젝트 디렉토리로 진입한 뒤 (본인의 경우 my_fastapi_app 디렉토리), 아래와 같이 명령어 입력.

uvicorn app.main:app --reload

성공하면 다음과 같이 진행된다.

4.1 POST

curl -X POST "http://127.0.0.1:8000/process/" -F "file=@/Users/kimyoungmin/my_fastapi_app/input_data/1329490_7069718_20151589.dcm"

실제 dicom 파일을 POST를 통해 입력한다.
성공한다면 위와 같은 표시가 뜰 것이다.
output_data 디렉토리에서 확인해보면,
결과가 제대로 나온 것을 확인할 수 있다.

4.2 GET

curl -X GET "http://127.0.0.1:8000/download_png/?file_path=/Users/kimyoungmin/my_fastapi_app
/output_data/1329490_7069718_20151589.png" -o /Users/kimyoungmin/Downloads/downloaded_image.png