이 프로젝트는 NVIDIA DeepStream SDK의 하드웨어 가속 디코딩과 ONNX Runtime을 결합하여, 대장 내시경 비디오에서 용종(Polyp)을 실시간으로 탐지하고 분석하는 AI 애플리케이션입니다.
- DeepStream 가속: NVIDIA GPU(NVDEC)를 활용한 고성능 비디오 디코딩 (GStreamer Pipeline)
- AI 추론: U-Net 기반의 Segmentation 모델 (ONNX)을 사용한 정밀 탐지
- 실시간 분석: 용종의 크기 측정, 위험도 분류, 실시간 FPS 모니터링
- 사용자 인터페이스: OpenCV 기반의 직관적인 GUI 및 키보드 제어 (일시정지, 저장 등)
본 프로젝트는 모델의 강건성(Robustness)을 확보하기 위해 성격이 다른 두 개의 공개 데이터셋을 활용하여 학습 및 검증을 수행하였습니다.
1. 학습 데이터 (Train Set): Kvasir-SEG
- 수량: 1,000장 (Images & Masks)
- 설명: 일반적인 용종 데이터 1,000장을 사용하여 모델이 용종의 보편적인 특징을 학습하도록 했습니다.
- 경로:
data/train/images,data/train/masks
2. 검증 데이터 (Validation Set): Kvasir-Sessile
- 수량: 196장 (Images & Masks)
- 설명: 탐지가 어려운 편평형(Sessile) 용종 데이터를 별도의 검증셋으로 구성하여, 학습하지 않은 어려운 케이스에 대한 일반화 성능을 평가했습니다.
- 경로:
data/val/images,data/val/masks
- OS: Windows 10/11 (WSL2 Ubuntu) 또는 Linux
- GPU: NVIDIA GeForce RTX 시리즈 (CUDA 지원 필수)
- Driver: 최신 NVIDIA 드라이버 설치 권장
- Docker: Docker Desktop (Windows) 또는 Docker Engine (Linux)
polyp_deepstream_app_v1.7.tar- Docker 이미지 파일kvasir/- 학습/검증 데이터셋 (Kvasir-SEG, Kvasir-Sessile)polyp_videos/- 테스트용 비디오 파일
이 시스템은 Docker 컨테이너 환경에서 동작합니다. 아래 순서대로 진행하면 모델 학습부터 실시간 비디오 추론까지 모든 과정을 재현할 수 있습니다.
Google Drive에서 파일을 다운로드한 후, 아래와 같은 폴더 구조를 만들어주세요.
MyProject/
├── polyp_deepstream_app_v1.7.tar ← [필수] Docker 이미지 파일
├── checkpoints/ ← (빈 폴더 생성) 모델이 저장될 곳
├── video/ ← [필수] 테스트용 .mp4 파일 위치
├── data/
│ ├── train/ ← Kvasir-SEG 데이터
│ │ ├── images/
│ │ └── masks/
│ └── val/ ← Kvasir-Sessile 데이터
│ ├── images/
│ └── masks/
├── train.py ← 학습 코드
├── validation.py ← 검증/시각화 코드
├── onnx_conversion.py ← ONNX 변환 코드
├── benchmark.py ← 성능 벤치마크 코드
├── test_polyp_detection.py ← Unit test 실행 코드
└── realtime_deepstream_gui.py ← 메인 실행 코드
WSL 터미널에서 프로젝트 폴더로 이동하여 이미지를 로드합니다.
# 폴더 이동
cd /mnt/c/Users/User/Downloads/MyProject
# 이미지 로드 (약 5~10분 소요)
docker load -i polyp_deepstream_app_v1.7.tar# [중요] HOST_DIR을 본인의 실제 폴더 경로로 수정하세요
HOST_DIR=/mnt/c/Users/User/Downloads/MyProject
docker run -it --rm --gpus all \
-e DISPLAY=$DISPLAY \
-v /tmp/.X11-unix:/tmp/.X11-unix \
-v ${HOST_DIR}:/workspace \
-v ${HOST_DIR}/data:/workspace/data \
-v ${HOST_DIR}/checkpoints:/workspace/checkpoints \
-v ${HOST_DIR}/video:/workspace/video \
--name polyp_pipeline \
polyp_deepstream_app:v1.7 \
/bin/bash컨테이너 내부(root@...:/workspace#)에서 아래 순서대로 실행하세요.
Step 1: 모델 학습 (train.py)
python3 train.py- Kvasir-SEG 데이터를 로드하여 U-Net 모델을 학습시킵니다.
- 결과:
checkpoints/best_metric_model_final.pth생성
Step 2: 단일 이미지 검증 (validation.py)
python3 validation.py- 학습된 모델이 잘 동작하는지 검증 데이터셋의 이미지를 추론해봅니다.
Step 3: ONNX 변환 (onnx_conversion.py)
python3 onnx_conversion.py- PyTorch 모델을 ONNX 포맷으로 변환합니다.
- 결과:
checkpoints/best_model_final.onnx생성
Step 4: 성능 벤치마크 (benchmark.py)
python3 benchmark.py- ONNX Runtime을 사용하여 검증 데이터셋 전체에 대한 처리 속도를 측정합니다.
Step 5: 실시간 비디오 탐지 (realtime_deepstream_gui.py)
python3 realtime_deepstream_gui.py- 비디오 파일에 대해 DeepStream 가속을 적용하여 실시간 탐지를 수행합니다.
realtime_deepstream_gui.py 실행 시 사용할 수 있는 단축키입니다.
| 키 (Key) | 기능 |
|---|---|
| SPACE | 영상 재생 / 일시 정지 |
| N | 다음 비디오 파일로 넘어가기 |
| B | 이전 비디오 파일로 돌아가기 |
| S | 현재 분석 화면 저장 (Screenshot) |
| ESC / Q | 프로그램 종료 |
-
학습 (Training):
DiceCELoss를 사용하여 클래스 불균형 문제를 해결하고,DiceMetric을 통해 최고 성능 모델을 저장합니다. -
변환 (Conversion): 실시간성을 위해
torch.onnx.export를 사용, 동적 배치(Dynamic Axes)를 지원하는 ONNX 모델로 변환합니다. -
추론 (Inference):
- Video Decoding: NVIDIA HW Decoder(NVDEC)를 사용하는 GStreamer Pipeline (
filesrc → qtdemux → h264parse → nvv4l2decoder) - AI Processing: ONNX Runtime (
CUDAExecutionProvider) - Visualization: OpenCV Drawing Functions
- Video Decoding: NVIDIA HW Decoder(NVDEC)를 사용하는 GStreamer Pipeline (
polyp_count_project/
├── README.md ← 프로젝트 설명서
├── .gitignore
└── polyp_count_project_code/
├── train.py ← 모델 학습 코드
├── validation.py ← 단일 이미지 검증 코드
├── onnx_conversion.py ← ONNX 변환 코드
├── benchmark.py ← 성능 벤치마크 코드
├── test_polyp_detection.py ← Unit test 실행 코드
├── realtime_deepstream_gui.py ← 메인 실행 코드 (GUI)
├── google_drive_link.txt ← 데이터/Docker 다운로드 링크
└── 프로젝트 구현 가이드.pdf ← 상세 가이드 문서
This project is for educational purposes.