An end-to-end TB segmentation system for chest X-rays

The problem

Tuberculosis often shows up first on a chest X-ray, but reading those films is slow and uneven, and a model that segments the affected regions has to be dependable, not just a notebook demo. We set out to take TB segmentation the whole way: from competing model architectures to a tracked, tuned, containerized service you can actually call.

What we built, and why

• Implemented two architectures, UNet and SegNet, in PyTorch and trained both on a single ml.g4dn.xlarge T4 GPU on SageMaker. UNet reached 0.9528 Dice and 0.9111 MIoU on the test set over 14 epochs with early stopping; SegNet reached 0.9569 Dice and 0.9180 MIoU from a much shorter run. Both land around 0.95 Dice, which is strong, and we are explicit that a fully matched head-to-head, equal epochs and budget, is the remaining step before ranking one over the other.
• Standardized preprocessing and augmentation for grayscale X-ray data, and kept training runs reproducible, so a result could be repeated and trusted rather than just rerun.
• Tracked every experiment in MLflow and drove hyperparameter search with Optuna, so model selection rested on recorded evidence, not guesswork. Each run's parameters, metrics, and artifacts are logged, which is what makes any UNet-versus-SegNet comparison honest in the first place.
• Engineered the path to production deliberately: exported the trained model to ONNX for portable, framework-independent inference, served it through a FastAPI layer in front of a SageMaker endpoint, and packaged the whole thing in Docker so it runs the same anywhere.
• Backed the system with a suite of 83 pytest tests across the data, model, and serving code, so regressions are caught before they reach the API.

The result

A working end-to-end TB segmentation system: two architectures trained to around 0.95 Dice on the test set, experiments tracked in MLflow, hyperparameters tuned with Optuna, the model exported to ONNX, and served through a FastAPI layer in front of a SageMaker endpoint, packaged with Docker. The 83-test suite is the evidence that it holds together rather than just running once.