의료 영상 처리를 위한 모듈형 PyTorch 프레임워크 MIP Candy

MIP Candy는 의료 영상 분할 파이프라인을 단일 메서드 구현만으로 구동할 수 있게 하면서, 레이어 교체, 교차 검증, ROI 기반 패치 추출, EMA, 실시간 메트릭 시각화 등 세밀한 제어를 지원하는 PyTorch 기반 모듈형 프레임워크이다.

저자: Tianhao Fu, Yucheng Chen

의료 영상 처리를 위한 모듈형 PyTorch 프레임워크 MIP Candy
본 논문은 의료 영상 분할 작업을 위한 새로운 PyTorch 기반 프레임워크인 MIP Candy(MIPCandy)를 소개한다. 의료 영상은 고해상도 3D 볼륨, 다양한 파일 포맷(NIfTI, DICOM, MHA), 불균형한 클래스 분포, 제한된 라벨 등 특수한 요구사항을 갖는다. 기존의 일반‑purpose 딥러닝 프레임워크(PyTorch, TensorFlow)는 이러한 요구를 직접 구현해야 하며, MONAI·TorchIO와 같은 컴포넌트 라이브러리는 조립 부담이 크고, nnU‑Net·MIST와 같은 엔드‑투‑엔드 파이프라인은 모듈성을 희생한다. MIP Candy는 이 두 극단 사이에 위치해, “단일 메서드 구현만으로 완전한 파이프라인을 구동”하면서도 각 구성 요소를 독립적으로 교체·재사용할 수 있게 설계되었다. **데이터 처리** mipcandy.data 모듈은 SimpleITK를 이용해 NIfTI·MetaImage·라스터 포맷을 자동 감지·로드하고, 필요 시 등방성 리샘플링 및 GPU 직접 배치를 지원한다. fast_save/fast_load는 safetensors 포맷을 사용해 제로‑카피 직렬화를 제공한다. 데이터셋은 torch.utils.data.Dataset을 상속한 베이스 클래스를 기반으로 하며, k‑fold 분할, 경로 내보내기, 디바이스 관리 기능을 포함한다. 주요 구현체로는 nnU‑Net 형식 데이터셋, 다클래스→이진 변환 래퍼인 BinarizedDataset, 데이터 병합 유틸리티가 있다. inspect() 함수는 각 케이스의 전경 바운딩 박스, 클래스 비율, 강도 통계 등을 자동으로 수집하고, 이를 기반으로 ROI‑shape를 추정한다. RandomROIDataset은 전경 오버샘플링(기본 33%)을 적용해 패치 기반 학습에 최적화된 샘플을 제공한다. **LayerT 구성 메커니즘** 전통적인 레이어 교체 방식은 다중 파라미터를 생성자에 넘기거나 서브클래싱을 요구한다. MIP Candy는 LayerT라는 경량 디스크립터를 도입해, 모듈 타입과 기본 생성 인자를 저장한다. assemble() 호출 시 실제 파라미터와 결합해 nn.Module 인스턴스를 만든다. 문자열 “in_ch”와 같은 지연 파라미터는 호출 시 전달된 값으로 자동 치환된다. ConvBlock2d, ConvBlock3d 등 대부분의 블록이 LayerT를 받아 컨볼루션·정규화·활성화 레이어를 자유롭게 교체할 수 있다. 이 설계는 2D/3D, 배치·그룹 정규화, ReLU·GELU 등 다양한 조합을 서브클래스 폭발 없이 지원한다. **학습 프레임워크** mipcandy.training 모듈은 Trainer 베이스 클래스와 TrainerToolbox 데이터클래스로 구성된다. Trainer는 학습 루프, 체크포인트 관리, 로그, 실시간 메트릭 저장 등을 담당하고, TrainerToolbox는 모델, 옵티마이저, 스케줄러, 손실, EMA 모델을 하나의 객체에 묶어 직렬화·복구를 단순화한다. SegmentationTrainer 프리셋은 다음을 기본값으로 제공한다. - 손실: Dice‑CrossEntropy (binary/multiclass 자동 선택) - 옵티마이저: SGD (momentum 0.99, Nesterov) - 스케줄러: polynomial LR decay - Gradient clipping, DeepSupervisionWrapper(옵션) - EMA(옵션) 학습 전 sanity check를 통해 출력 형태와 파라미터 수를 검증하고, 매 epoch마다 CSV 메트릭, 그래프, 최악 사례 예측 프리뷰를 저장한다. 학습 중단 시 epoch 단위 체크포인트를 복구해 연속 학습이 가능하다. 또한 검증 곡선에 유리 함수(quotient regression)를 피팅해 최대 예상 점수와 최적 종료 epoch, 남은 시간 등을 예측한다. **추론·평가·프론트엔드** mipcandy.inference는 Predictor 베이스 클래스를 제공해 파일‑레벨 예측·내보내기를 지원한다. 평가 모듈은 Evaluator와 EvalResult를 통해 per‑case 및 aggregate Dice, IoU 등 메트릭을 계산한다. metrics 서브패키지는 binary_dice, dice_similarity_coefficient, soft_dice 등을 포함한다. frontend 모듈은 Weights & Biases, Notion, MLflow와의 연동을 추상화해, 실험 로그, 파라미터, 아티팩트 등을 다중 플랫폼에 동시 기록할 수 있다. **생태계·확장성** MIP Candy는 bundle 개념을 도입해 모델·트레이너·프리딕터를 하나의 패키지로 묶어 배포한다. 각 번들은 core 프레임워크와 동일한 인터페이스를 따르므로, 새로운 아키텍처를 추가할 때 코어를 수정할 필요가 없다. 이는 MONAI의 “bundle”과 유사하지만, 전체 파이프라인을 자동으로 구성하는 nnU‑Net과 달리 사용자가 원하는 부분만 선택해 조합할 수 있다. **비교·위치** 표 1과 2에서 기존 프레임워크와 MIP Candy의 기능을 비교한다. MIP Candy는 완전한 파이프라인 제공과 모듈성 모두를 만족한다. 특히 실시간 메트릭 시각화, 검증 점수 예측, 다중 프론트엔드 추적 등 연구 투명성을 높이는 기능이 독보적이다. **결론** MIP Candy는 의료 영상 분할을 위한 현대적 요구를 충족하는 모듈형·완전형 프레임워크로, 레이어 교체를 위한 LayerT, 데이터셋 자동 ROI 추출, 깊은 감독, EMA, 검증 점수 예측 등 풍부한 기능을 제공한다. Python 3.12 이상의 최신 언어 기능을 활용해 가독성과 정적 검증을 강화했으며, Apache‑2.0 라이선스로 오픈소스 배포돼 연구 재현성과 커뮤니티 확장을 기대한다.

원본 논문

고화질 논문을 불러오는 중입니다...

댓글 및 학술 토론

Loading comments...

의견 남기기