DeepNeuro: Real-time Multi-modal Brain Tumor Segmentation using nnU-Net

The primary challenge in medical image segmentation is the massive class imbalance: brain tumor regions typically occupy less than 2% of the total voxel space in an MRI volume. Standard cross-entropy loss functions fail in this environment, as neural networks tend to converge towards predicting the background class, ignoring the tumor sub-regions.

Furthermore, medical scans are multi-parametric, requiring a combined analysis of four distinct modalities: FLAIR, T1, T1ce, and T2. Developing a robust engineering pipeline that can ingest 3D NIfTI files, align multiple channels, extract representative axial slices, Z-score normalize the voxel intensity distributions, and run inference within interactive latency limits on standard CPU web servers was a significant full-stack architectural challenge.

We architected a custom 2D nnU-Net style architecture in PyTorch. The model features 5 downsampling levels, dynamic padding for arbitrary axial sizes, and caps the maximum feature maps at 512 channels. By replacing standard BatchNorm layers with InstanceNorm2d, the network normalizes each slice and modality independently, maintaining stable gradients even at tiny batch sizes during training.

To address class imbalance, we implemented a combined Dice + Cross-Entropy Loss, which optimizes both voxel-level classification and region overlap. The preprocessing engine leverages NiBabel for robust NIfTI extraction, automatically detects and stacks sibling modalities in the folder structure, Z-score normalizes modalities individually, and falls back gracefully to single-channel duplication if modality files are missing. The model runs locally on a Flask backend, using CPU-optimized PyTorch weights to deliver real-time results.

DeepNeuro leverages a highly optimized model architecture inspired by the nnU-Net framework, which systematically outperforms vanilla U-Net on medical datasets. A major design decision was capping the maximum feature maps at 512 channels in the bottleneck. While a standard U-Net expands to 1024 channels, our custom nnU-Net style caps channels at 512. This parameter-efficiency decision reduced the model size from 31 million to 13.2 million trainable parameters (a 57.2% reduction), mitigating overfitting on slice extractions and speeding up inference. Training was executed on a Kaggle Tesla T4 GPU. The dataset consists of all 1,251 patients from the official BraTS 2021 dataset. During training, we applied joint, synchronized spatial augmentations (random horizontal/vertical flips and orthogonal 90/180/270-degree rotations) to both the 4-channel inputs and ground truth masks to ensure spatial alignment. The model converged over 5 epochs with the AdamW optimizer (learning rate 2e-4), successfully reducing the average loss from 2.79 to 1.08. On the backend, the inference pipeline is designed to be highly adaptive. If a user uploads a single modality slice, the engine duplicates the channels to match the 4-channel model input. If a NIfTI volume is uploaded, it scans the directory for counterpart modality names, stacks them, extracts the central axial slice, and outputs a clean, colorful RGB overlay mapping Necrotic Tumor Core (NCR - Red), Peritumoral Edematous Tissue (ED - Green), and Enhancing Tumor (ET - Blue). Additionally, a synthetic demo data generator was built to create randomized mathematical brain tumor clusters, letting users evaluate the pipeline with one click without uploading heavy NIfTI files.

The result is a production-ready medical imaging prototype. By capping feature maps and leveraging InstanceNorm, the custom nnU-Net style model achieved an average CPU latency of 181.96 milliseconds per slice, compared to 768.70 milliseconds for the standard U-Net. This represents a 4.2x latency speedup and increases the model's throughput from 1.30 slices/second to 5.50 slices/second, making it viable for real-time web deployment.

Building DeepNeuro allowed us to gain hands-on experience with medical image processing standards, advanced spatial data augmentation, and multi-modal neural networks. The optimization techniques implemented—ranging from custom loss formulations to model compression and modality auto-alignment—demonstrate senior-level deep learning engineering and clean backend design.