Introduction: Chest X-Rays serve as a pivotal diagnostic tool globally, enabling physicians to assess and diagnose various illnesses cost-effectively. However, leveraging these images for diagnosis presents challenges, primarily due to the scarcity of datasets featuring associated doctor's notes, hindering the development of accurate computer-aided diagnosis models. In this project, we utilized the NIH Chest X-rays dataset comprising 112,120 X-ray images with disease labels derived from 30,805 unique patients. These labels were generated through Natural Language Processing, extracting disease classifications from radiological reports linked to each X-Ray image. Additionally, the dataset includes a limited number of disease region bounding boxes (880 bounding boxes), delineating the areas within the X-Ray indicating the presence of disease. During our Exploratory Data Analysis (EDA), we identified patients with recorded ages exceeding 120 years. Given the implausibility of such ages, we inferred inaccuracies in the data. Below are two graphs illustrating the dataset's distribution after correcting these erroneous ages, as well as the distribution of diseased and non-diseased X-Ray images. Leveraging the dataset, we developed three distinct models aimed at classifying both the presence and localization of disease within Chest X-Ray images. The first model concentrated solely on binary classification, discerning the presence or absence of disease in X-Ray images. The second model focused on generating bounding box coordinates, pinpointing specific areas within X-ray images indicating the presence of disease. This approach utilized the limited number of bounding boxes provided in the original dataset. The final model centered on creating new bounding boxes using information provided by the entire X-Ray image dataset. This was achieved through transfer learning- the trained weights of the classification model were used to initialize half of the layers in this new segmentation model, and the latter half were trained on the bounding box information.

Methodology: Model 1: Binary Classification The binary classification model was constructed using a Convolutional Neural Network (CNN) architecture, designed to ingest both the X-Ray image and tabular patient data as inputs. The CNN architecture comprises four convolutional layers integrating L2 regularization, batch normalization, max pooling, and dropout mechanisms. Additionally, the network incorporates a dense layer with ReLU activation for training the patient data. The output from the final convolutional layer is flattened and combined with the output from the dense patient data layer. Subsequently, this concatenated representation is fed into a final dense layer with sigmoid activation, forming the binary classifier.

Model 2 and 3: Baseline Segmentation and Segmentation with Transfer Learning The segmentation model is employed to delineate bounding boxes onto a unified segmentation mask, particularly when multiple boxes are present within an X-Ray image. Each bounding box is characterized by four coordinates [x, y, w, h], denoting the top-left corner coordinates (x, y) along with the width (w) and height (h) of the box. The model architecture utilizes a U-Net framework, which takes X-Ray images and patient data as inputs. The downsampling path of the U-Net comprises four convolutional layers, each followed by max pooling, alongside a dense layer for incorporating tabular patient data. Upon completion of the downsampling process, the attributes are concatenated at a bottleneck layer before proceeding to the upsampling path. The upsampling path consists of four convolutional layers, each accepting a concatenation of the preceding layer with the corresponding downsampling output. The final layer of the upsampling path incorporates three convolutional layers with group normalization and sigmoid activation, yielding the segmentation output. The output of the segmentation model is a single array, matching the size of the input image (1024 x 1024). Each array entry represents a pixel in the image, denoted by either "yes" or "no." A "yes" entry signifies the pixel's predicted association with the diseased region of the image. Both the baseline and transfer learning segmentation models share identical structures; however, the distinction lies in how the model initializes its weights. The baseline segmentation model begins with uninitialized weights and is exclusively trained using the 880 X-Ray images with corresponding bounding boxes. Conversely, the transfer learning model initializes its weights from the CNN trained on the entire dataset of 112,120 X-Ray images and subsequently fine-tunes on the 880 segmentation samples.

Results: Model 1: Binary Classification Our CNN underwent two rounds of training: initially on a subset of 12,000 images along with corresponding patient data and disease labels, and subsequently on the full dataset. The purpose of the initial run was to expedite CNN training due to the large dataset size. However, the achieved accuracy scores were suboptimal, with a training accuracy of 0.6171 and a testing accuracy of 0.5730. Upon refining our data loading process to expedite training, we conducted a second round of training on the complete dataset. As anticipated, this resulted in improved accuracy, with a training accuracy of 0.6366 and a test accuracy of 0.5804. We attribute the relatively low accuracy scores to our simplification of the classification model from multi-class to binary. Originally, the dataset encompassed thirteen diseases along with a "not diseased" designation for labeling each X-Ray image. By amalgamating the thirteen diseases into a single "diseased" label, we may have compromised the CNN's ability to discern nuanced attributes in each X-Ray image causing the model to predict diseased or not diseased at a lower accuracy then expected.

Model 2: Baseline Segmentation and Segmentation with Transfer Learning Our baseline segmentation model was trained on the 880 bounding box regions and associated X-Ray images that were originally sourced from the NIH Chest X-Ray dataset. Evaluation of this model was conducted using both accuracy metrics and a measure of segmentation quality known as Intersection over Union (IoU) Positive. This metric assesses the overlap between predicted and true bounding boxes by dividing their intersection area by their union area. The baseline segmentation model achieved a training accuracy of 0.8429, a testing accuracy of 0.8289, a training IoU Positive of 0.2373, and a testing IoU Positive of 0.1969. These results were noteworthy, particularly given the limited amount of training data available. The segmentation model with transfer learning was trained on the smaller segmentation dataset while using the initialized weights from the CNN model. The learning rate for transferred weights was set to 1e-6, and the rate for randomly initialized layers to 1e-3. Evaluation of this model employed the same metrics as the baseline. It exhibited a training accuracy of 0.8686, a testing accuracy of 0.8557, a training IoU Positive of 0.2639, and a testing IoU Positive of 0.2071. The hypothesis driving this model's development was its anticipated performance improvement due to the larger training dataset. However, uncertainties existed regarding the efficacy of transfer learning in generating additional bounding boxes for X-Rays lacking them. As we expected, the segmentation model with transfer learning surpassed our baseline model across all evaluation metrics. This success underscores the effectiveness of transfer learning despite initial uncertainties regarding its application to X-Rays lacking associated bounding boxes.

Challenges: The primary obstacle encountered during this project revolved around processing time constraints and the limitations imposed by the available computational resources, largely attributable to the dataset's size. We undertook various strategies to expedite training and augment computational capacity, including reducing the size of the training set, streamlining the CNN architecture by minimizing convolutional layers, limiting the number of training epochs, leveraging Tensorflow Metal, and maximizing GPU allocation on platforms like Google Colab. Our pivotal breakthrough occurred when we identified data loading as the bottleneck in the training pipeline. To address this, we devised a solution involving the aggregation of all multi-modal data and partitioning the dataset into sizable shards. With this approach, we divided the 112,120 X-Ray images into 40 .tfrec files using the TFRecord format. Each shard contained a few thousand images paired with their respective multi-modal data, encompassing patient demographics and clinical findings. This strategic partitioning enabled the model to ingest thousands of images simultaneously during training, as opposed to processing them sequentially, significantly enhancing training efficiency. Further challenges were encountered during the implementation of the transfer learning model, necessitating meticulous planning and collaborative efforts across all models to ensure appropriate architecture and desired outcomes. Consequently, each model had to be executed sequentially rather than simultaneously to guarantee consistency in architectures.

Reflection: Our model largely performed as expected, with the success of the segmentation model utilizing transfer learning marking a significant milestone in understanding the application of transfer learning within this field of study. However, the margin of performance increase observed by the transfer learning approach was quite small. This may well be due to the fact that the classification model itself was not entirely performance either. We believe that this comes from the difficulty of the task at hand, rather than any specific design failing. Ultimately, we aimed to run a more robust CNN that incorporated multi-class classification (our stretch goal), which we believe would have increased the accuracy of all of our models. This goal was made challenging by the fact that many X-Ray images had multiple diseases present within the image. This complexity necessitates future research to develop a CNN capable of accurately identifying multiple diseases per image. Throughout the project, our approach to model training evolved as we grappled with the time-intensive nature of training on the full dataset. Initially, we pivoted to training the CNN on a subset of the data to expedite results, albeit even this smaller subset required considerable time. It wasn't until later in the process that we resolved data loading issues and could train on the entire dataset. Consequently, time constraints limited parameter tuning and model architecture adjustments, which could have significantly enhanced our results and predictive capabilities given more time. Our key learnings underscored the critical importance of data engineering in data science projects, highlighting the need to address data-related challenges early in the process. Additionally, we learned how to implement segmentation and transfer learning, emphasizing their potency as deep learning techniques capable of substantially improving model predictive performance, providing valuable insights for future projects.

References:

National Institutes of Health Chest X-Ray Dataset. “NIH Chest X-Rays.” Kaggle, 21 Feb. 2018, www.kaggle.com/datasets/nih-chest-xrays/data?resource=download&select=images_011. “TPU-Speed Data Pipelines: Tf.Data.Dataset and Tfrecords | Google Codelabs.” Google, Google, codelabs.developers.google.com/codelabs/keras-flowers-data#4. Accessed 20 Apr. 2024. Tzeng, I-Shiang. “A Review of Recent Advances in Deep Learning Models for Chest Disease Detection Using Radiography.” National Library of Medicine, 3 Jan. 2023, www.ncbi.nlm.nih.gov/pmc/articles/PMC9818166/. Von Csefalvay, Chris. “Transfer Learning: The DOS and Don’ts.” Medium, HCLTech-Starschema Blog, 18 Mar. 2019, medium.com/starschema-blog/transfer-learning-the-dos-and-donts-165729d66625. Y. Wu and K. He, "Group normalization" in arXiv:1803.08494, 2018

GitHub Repository: https://github.com/BrownDSI/DL_project/blob/master/main.ipynb