Understanding the tremendous potential of AI to cut physicians’ workloads and reduce time to diagnosis, AltexSoft invested effort in creating a prototype of a decision support tool for chest X-ray analysis. It aimed to perform three machine learning tasks: lung segmentation to identify the boundaries of lungs; lung disease classification to predict a probability of pneumonia, fibrosis, and pneumothorax; and pneumothorax (collapsed lung) mask creation to localize this dangerous condition if detected.
The tool targets both medical experts and patients, serving as a source of a second opinion about the presence of lung abnormalities. If required, the functionality can be relatively quickly extended to more disorders and non-X-ray medical imaging, primarily CT scans.
Medical imaging datasets are extremely hard to obtain because of strict privacy regulations in the healthcare sector and the high cost of manual annotation as each instance must be examined by three professionals. Creating them from the ground up exceeds the time and budget constraints of the prototype project. Instead, our team explored the resources of Kaggle, the world’s largest community of data scientists, boasting 50,000 public datasets. After doing some research, we selected, cleaned, and pre-processed three datasets with the greatest relevance to our goals. This included de-identified chest x-rays and corresponding lung masks collected for a tuberculosis control program; over 112, 000 radiographs with disease labels from about 30,000 unique patients; and a dataset provided by the Society for Imaging Informatics in Medicine for pneumothorax segmentation.
Our set of tasks belongs to the computer vision field where convolutional neural networks (CNNs) shine. After running numerous experiments, we defined CNN architectures that produced the best results for each of the three problems we outlined. In our tool, lung segmentation rests on U-Net architecture with B5 EfficientNet backbone (feature extractor). The model achieves 0.93 of IoU (Intersection-over-Union), a metric that shows the area of overlap between the predicted segmentation and the ground truth. Disease classification is held by B4 EfficientNet. And pneumothorax localization is based on Feature Pyramid Network (FPN) with B5 EfficientNet backbone. It hits a Dice score of 0.83 (Dice score measures the similarity of two samples.)
All three CNNs were united into a single tool with a simple interface that would be comprehensible for everybody, no matter the level of medical expertise. The software enables uploading X-rays to be processed by the AI-based engine. The output is the image with the lung mask that shows the borders of the organ. If there is a pneumothorax, the damaged area is marked with red coloring. Also, the tool calculates the probability level of three lung conditions — pneumothorax, pneumonia, and fibrosis. The user can export results in one click to share them with a doctor or colleague.
The prototype was completed in 2 months by a team of 4 experts: 2 ML engineers, 1 software engineer, and 1 UX/UI designer. The total human effort amounted to 4-5 person/month.
The technology stack comprised Python, the aiohttp package, Docker, React, and JavaScript. As for the machine learning stack, it included U-Net, EfficientNet, Feature Pyramid Network (FPN), PyTorch, Albumentations, and OpenCV.