PR3DICTR: A modular AI framework for medical 3D image-based detection and outcome prediction

1. 문제 정의 (Problem)

최근 3 차원 (3D) 의료 영상 데이터와 딥러닝 기반의 컴퓨터 지원 의사결정이 의료 분야에서 급격히 중요해지고 있습니다. 그러나 이러한 기술의 발전과 적용에는 다음과 같은 주요 장애물이 존재합니다.

• 개발의 복잡성과 비표준화: 3D 의료 영상 기반의 딥러닝 모델 개발은 상당한 계산 자원과 기술적 전문성을 요구합니다. 연구 그룹 내외부에서 워크플로우가 다양하고 비표준화되어 있어, 결과의 재현성 (Reproducibility) 이 낮고 개발 부담이 큽니다.
• 기존 프레임워크의 한계:
  • PyTorch/MONAI: 유연성이 높지만, 의료 영상 특화 작업을 위해 많은 코드를 직접 구현해야 하는 번거로움이 있습니다.
  • No-code 도구 (예: Ludwig): 설정 파일로 모델을 정의할 수 있으나, 특정 의료 영상 분류 작업에 적용하기에는 너무 광범위하거나 개별 예측 모델의 요구사항에 적응하기 어렵습니다.
• 필요성: 3D 의료 영상 데이터에 기반한 분류 모델링을 위해 단순성, 모듈성, 표준화를 균형 있게 갖춘 프레임워크가 필요합니다.

2. 방법론 (Methodology)

저자들은 PR3DICTR(Platform for Research in 3D Image Classification and sTandardised tRaining) 이라는 오픈소스 프레임워크를 제안했습니다. 이 프레임워크는 PyTorch 와 MONAI 를 기반으로 구축되었으며, 다음과 같은 핵심 아키텍처와 워크플로우를 제공합니다.

A. 핵심 워크플로우

1. 데이터 준비 (Data Curation): 원시 DICOM 데이터를 구조화된 NumPy 형식으로 변환하고, 환자 ID 기반의 디렉토리 구조로 정리합니다. (프레임워크는 이 전 단계가 완료된 데이터를 가정합니다.)
2. 설정 파일 (Configuration): 모든 모델 정의, 하이퍼파라미터, 학습 설정은 YAML 기반의 설정 파일 (config) 로 관리됩니다. 이를 통해 사용자는 복잡한 코드 수정 없이 두 줄의 코드만으로 실험을 시작할 수 있습니다.
3. 모델 아키텍처:
   • 이미지 인코더 (Image Encoder): 3D 영상 데이터 (CT, PET, MRI 등) 에서 특징을 추출합니다. ResNet, DenseNet, EfficientNet, ConvNeXt, ViT 등 다양한 CNN 및 Transformer 아키텍처를 지원합니다.
   • 출력 모듈 (Output Module): 추출된 영상 특징과 임상 데이터 (Tabular data) 를 통합합니다. 다중 레이블 분류를 지원하며, 임상 데이터의 결합 위치와 레이어 구조를 유연하게 설정할 수 있습니다.
   • MLP 지원: 영상 데이터 없이 임상 데이터만 사용할 경우, 다층 퍼셉트론 (MLP) 으로 자동 전환됩니다.
4. 학습 및 최적화:
   • 데이터 로딩 및 증강: MONAI 의 다양한 데이터셋 인터페이스 (Cache, SmartCache 등) 를 지원하며, 결정적 변환 (크롭, 정규화) 과 비결정적 변환 (랜덤 회전, MixUp 등) 을 제공합니다.
   • 하이퍼파라미터 최적화: Optuna를 통합하여 자동화된 하이퍼파라미터 탐색을 수행하며, K-폴드 교차검증을 지원합니다.
   • 평가: 이진 분류 및 생존 분석 (Time-to-event) 을 위한 다양한 지표 (AUC, C-index, Brier score 등) 와 시각화 도구 (칼리브레이션 플롯, ROC 곡선 등) 를 내장하고 있습니다.
   • 추적: Weights & Biases (W&B) 와 연동되어 실험 추적 및 실시간 모니터링이 가능합니다.

B. 모듈성 (Modularity)

프레임워크는 데이터 로딩, 모델 아키텍처, 학습, 평가 등 각 구성 요소를 독립적으로 교체하거나 확장할 수 있도록 설계되었습니다. 사용자는 설정 파일에서 모듈을 선택하거나 커스텀 모듈을 '플러그인' 방식으로 추가할 수 있습니다.

3. 주요 기여 (Key Contributions)

• 표준화된 3D 의료 영상 분류 프레임워크: 3D 의료 영상 기반 예측 모델 개발을 위한 오픈소스 플랫폼을 최초로 제공하여, 연구 간 워크플로우의 일관성을 확보했습니다.
• 높은 접근성과 유연성의 균형: 설정 파일 (config) 기반의 간소화된 인터페이스를 제공하면서도, 내부 모듈을 자유롭게 수정할 수 있는 높은 유연성을 유지합니다.
• 다중 모달리티 및 데이터 통합: CT, PET, 방사선 선량 분포, segmentation 마스크 등 다양한 3D 영상 모달리티와 임상 표본 데이터 (Tabular data) 를 통합하여 학습할 수 있는 기능을 제공합니다.
• 재현성 강화: 모든 실험 설정, 모델 가중치, 평가 결과를 체계적으로 저장하고 로깅하여, 연구의 투명성과 재현성을 극대화합니다.

4. 결과 (Results)

• 실증 사례 (NSCLC-Radiomics 데이터셋): The Cancer Imaging Archive 의 NSCLC-Radiomics 데이터셋을 사용하여 성별 분류 (Sex classification) 태스크에 프레임워크를 적용했습니다.
• 성능: 기본 제공되는 ResNet-10 모델을 사용하여 테스트 세트에서 거의 완벽한 성별 구분 능력을 보였으며, 모델의 보정 (Calibration) 도 양호했습니다.
• 코드 효율성: 데이터 전처리와 모델 학습을 위해 Jupyter Notebook 예제를 제공하며, 실제 모델 실행은 두 줄의 코드만으로 가능함을 시연했습니다.

5. 의의 및 중요성 (Significance)

• 연구 효율성 증대: 기존에 연구 그룹 내에서 개별적으로 구현되던 DL 모델 개발 코드를 통합함으로써, 구현에 따른 변이를 줄이고 다양한 질병 부위 (두경부, 폐암 등) 간 모델 비교를 용이하게 했습니다.
• 표준화 및 재현성: 의료 AI 연구의 핵심 과제인 재현성 문제를 해결하기 위해, 설정 파일과 자동화된 평가 시스템을 통해 방법론적 선택 사항을 명확히 문서화하고 추적 가능하게 만들었습니다.
• 확장성: 현재 이진 분류와 생존 분석에 초점을 맞추고 있지만, 모듈형 설계 덕분에 다중 클래스 분류, 불확실성 정량화 (Uncertainty Quantification), 주석 맵 (Attention maps) 등 향후 기능 확장이 용이합니다.
• 미래 지향성: 의료 AI 의 임상 적용을 가속화하기 위해, 복잡한 개발 과정을 단순화하면서도 투명성과 유연성을 유지하는 이상적인 프레임워크로 자리매김할 것으로 기대됩니다.

결론적으로, PR3DICTR 은 3D 의료 영상 기반 딥러닝 모델 개발의 진입 장벽을 낮추고, 연구의 표준화와 재현성을 동시에 달성하기 위해 설계된 강력하고 유연한 도구입니다.