Capsule Networks for 3D Medical Imaging Segmentation in Julia
Julia's medical imaging stack currently relies on 3D U-Net variants for segmentation, which lose spatial relationships through max-pooling. This project implements 3D capsule network architectures in Lux.jl that preserve part-whole spatial hierarchies through dynamic routing-by-agreement. I will build reusable capsule layer primitives (PrimaryCapsule3D, ConvCapsule3D, routing modules), two full architectures (3D SegCaps and a SegCaps-UNet hybrid), and custom GPU-accelerated routing kernels using KernelAbstractions.jl. The full pipeline integrates with MedPipe3D.jl for data loading, training, and evaluation. Deliverables: (1) Capsule layer primitives as Lux.jl modules, (2) 3D SegCaps architecture, (3) SegCaps-UNet hybrid, (4) GPU routing kernels, (5) training pipeline with MedPipe3D.jl, (6) benchmarks across all 10 Medical Segmentation Decathlon tasks (Dice, HD95, cross-task transfer vs 3D U-Net), (7) documentation and pretrained weights, (8) PRs to JuliaHealth.
Project details
Technologies
Not listed in the archive