This is the repo of the Applied Machine Learning project by Giovanni Battista Esposito
The goal of the project is to build a multi-label classification model for the X-Ray Chest dataset.
The project is developed in the following stages: first, there is a data exploration part, in which the statistical propetries of the dataset will be analyzed. Second, there is the design of a multimodal model combining tabular data with images; finally there is the training, calibration and testing of the model.
The dataset is structured as follows:
<class 'pandas.core.frame.DataFrame'>
RangeIndex: 108494 entries, 0 to 108493
Data columns (total 21 columns):
# Column Non-Null Count Dtype Description
--- ------ -------------- ----- -----------
0 Image_name 108494 non-null object Filename of the x-ray image
1 Patient_ID 108494 non-null int64 Anonymized ID of the patient
2 Study 108494 non-null int64 Study group
3 Sex 93829 non-null object Patient's sex
4 Age 93829 non-null float64 Patient's age
5 ViewCategory 108494 non-null object Type of the X-ray image (frontal/lateral)
6 ViewPosition 108494 non-null object View position of the X-ray (PA, AP, Lateral, LL)
7 Atelectasis 108494 non-null int64 Disease
8 Cardiomegaly 108494 non-null int64 Disease
9 Consolidation 108494 non-null int64 Disease
10 Edema 108494 non-null int64 Disease
11 Enlarged Cardiomediastinum 108494 non-null int64 Disease
12 Fracture 108494 non-null int64 Disease
13 Lung Lesion 108494 non-null int64 Disease
14 Lung Opacity 108494 non-null int64 Disease
15 No Finding 108494 non-null int64 No Disease
16 Pleural Effusion 108494 non-null int64 Disease
17 Pleural Other 108494 non-null int64 Disease
18 Pneumonia 108494 non-null int64 Disease
19 Pneumothorax 108494 non-null int64 Disease
20 Support Devices 108494 non-null int64 Disease
Each observation has the following components:
- A tabular dataset (Sex, Age, View Position, View Category) with information about the patient and the scan
- An image of variable size and resolution that is the x-ray scan itself
- A collection of 14 boolean labels indicating diseases affecting the patient
All the diseases ar enot mutually exclusive, meaning that a patient can be affected by multiple conditions, exept for "No Finding" that indicates the absence of any disease.
Further information about the dataset can be found in the source repository
Note that I only used train2 images directory. Although the original source included additional images for the testing set, they were missing the tabular features (like sex, id, category, etc.). In order to build a complete model that handles both, I decided to only use data under train2 directory, which is enriched by those information.
The first stage of the project involves an explorative analysis on the tabular dataset
With the exploration, some considerations on the data emerged:
There are approximately ~13% of the data with missing both sex and age. Observations that miss sex always miss age and vice versa. Two possible strategies have been considered: dropping rows with missing sex/age, and imputing. Both the strategies are likely to introduce bias in the model to the goal is to choose the one that will affect less the model's performances First, I visually analyzed if the distribution of the labels of incomplete rows is similar to the complete rows' one. A similar distribution would mean that missing observations are Missing Completely At Random, meaning that the missingness does not depend on the observations themselves but is random.
It is plain to see that the distribution of the classes changes a lot, meaning that data is not MCAR but could be Missing At Random or Missing Not At Random. Theory suggests that in this situation is not safe to drop missing rows, so I decided for the imputation way.
To impute an observation's sex and age, first of all I checked if there were other observations from the same patient in which this data was present, but there were none. So, for Sex I decided to create a third category alongside Male and Female to indicate the uncertainty. For Age, the strategy was a median imputation
Sex, ViewPosition and ViewCategory are categorical variables, meaning there is the need of an encoding. Sex has 3 categories (Male, Female, Missing), ViewCategory has 2 (Frontal, Lateral) and ViewPosition has 4 (AP, PA, LL, Lateral). They were one-hot encoded, ending up with 6 boolean variables. Study is also a categorical variable with over 180 values. I decided to drop it as the study number has no numerical meaning and a 180-categories one-hot encoding would have complicated a lot the final model architecture.
Alongside them, Age is a numerical variable. It has been scaled in a range [0, 1].
As introduced before, the variable NoFinding is a boolean variable being True if all the other disease variables are False, False otherwise.
I believe that in a multi-label classification model, this is an overcomplication and it is not useful. Including it would require the introduction of a constraint in the training loss, forcing the logits to satisfy
The splitting strategy for the dataset was not a simple random split. In fact, the problem is that dividing images of the same patient in different sets, or, for example, a frontal image in a different set of the same, lateral image, would have been spoiled data leakage into the model. So the splitting was done by unique Patient ID's, to keep the intergity of the model. So, unique IDs (and so the tabular data and the images) were splitted in 4 different sets:
- 60% of the patients in the training set
- 20% of the patients in the validation set
- 10% of the patients in the calibration set
- 10% of the patients in the testing set
Alongside the tabular info, there are corresponding images. Visualizing a sample of the images:
The images included in the dataset are grayscale of variable width and height. They will be transformed to a 224x224, RGB format to leverage on a pretrained ResNet for their encoding.
When dealing with images, it is also useful to do some processing and augmentations. The following transformations have been applied to the training set:
- Contrast Enhancement (CLAHE)
- Resize to 224x224
- Random Rotation (up to 10°)
- Random Horizontal Flip (50% probability)
- Random Crop (up to 10% of the original image)
- Color Jitter (brightness and contrast variations)
- Conversion to a 3-channels format
- Scaling in [0, 1] range
- Normalization of the channels around ImageNet dataset values
Other sets (validation, calibration, testing) are only subject to contrast enhancement, resizing, conversion to RGB, scaling in [0,1] and normalization.
This is an example of augmented images (different from the precious picture):
The model is built with 2 different encodings for tabular and image data.
- The Tabular Encoder is a Fully Connected Network with an input layer, an activation function and an output layer. It takes as input the 7 tabular features and maps them into a 32-dimensional hidden space
- The Image Encoder is built upon a pretrained ResNet18 model mapping each image into a 512-dimensional latent space. At the bottom of the ResNet has been added a linear projection layer to also map them into a 32-dimensional space
- The Fusion Encoder is again a FCN that takes the 32-dimensional tabular vector and the 32-dimensional image vector, fuse them through two hidden linear layers and outputs a 32-dimensional vector
- The Output Layer is a final, linear layer returning ta 13-dimansional vector of each class' logit value
====================================================================================================
Layer (type:depth-idx) Output Shape Param #
====================================================================================================
Model [1, 13] --
├─TabularEncoder: 1-1 [1, 32] --
│ └─Linear: 2-1 [1, 32] 256
│ └─ReLU: 2-2 [1, 32] --
│ └─Linear: 2-3 [1, 32] 1,056
├─ImageEncoder: 1-2 [1, 32] --
│ └─ResNet: 2-4 [1, 32] --
│ │ └─Conv2d: 3-1 [1, 64, 112, 112] (9,408)
│ │ └─BatchNorm2d: 3-2 [1, 64, 112, 112] (128)
│ │ └─ReLU: 3-3 [1, 64, 112, 112] --
│ │ └─MaxPool2d: 3-4 [1, 64, 56, 56] --
│ │ └─Sequential: 3-5 [1, 64, 56, 56] (147,968)
│ │ └─Sequential: 3-6 [1, 128, 28, 28] (525,568)
│ │ └─Sequential: 3-7 [1, 256, 14, 14] (2,099,712)
│ │ └─Sequential: 3-8 [1, 512, 7, 7] (8,393,728)
│ │ └─AdaptiveAvgPool2d: 3-9 [1, 512, 1, 1] --
│ │ └─Linear: 3-10 [1, 32] 16,416
├─FusionLayer: 1-3 [1, 32] --
│ └─Linear: 2-5 [1, 32] 2,080
│ └─ReLU: 2-6 [1, 32] --
│ └─Linear: 2-7 [1, 32] 1,056
├─Linear: 1-4 [1, 13] 429
====================================================================================================
Total params: 11,197,805
Trainable params: 21,293
Non-trainable params: 11,176,512
Total mult-adds (Units.GIGABYTES): 1.81
====================================================================================================
Input size (MB): 0.60
Forward/backward pass size (MB): 39.74
Params size (MB): 44.79
Estimated Total Size (MB): 85.13
====================================================================================================
The code of the model class is available in src/model.py
The model was trained with a Weighted Binary Cross Entropy Loss, to handle class imbalance. The class weights were computed on the training set as $$ w_c = \frac{N}{C \times \sum_{i=1}^Nc_i} $$
Namely, the weight of a class
The training procedure used an AdamW optimizer with dynamic learning rate; to control it, I introduced a policy to reduce the learning rate when the validation loss stops its descending. Furthermore, the training used a gradient scaling to introduce mixed precision and optimize better the procedure, as the huge expense that it takes to train a model on over 100GB of data.
In the training script there is the full code used to train the model.
Given the huge dimensionality of the data, the training of the model required approximately 17 hours (logs are here) for 15 epochs.
From the loss curves, it appears that the model is stil in its underfitting stage. I believe the reason stands (i) in a not sufficient number of epochs and (ii) in the amount of dropout introduced to enhance generalization. But I observed with some hyperparameter tuning that this amount of droput is necessary in the inference stage.
When dealing with logits and multi-class models with imbalance, there is the risk that the predicted probabilities are not well calibrated. Calibration of a model can be checked though the reliability diagrams. Ideally, if all the classes are well-calibrated, their curves should be a perfect diagonal line, meaning that it is predicting the right probabilities and they are coherent with the reality.
I first reintroduced No Finding to check also its calibration, as
The model is generally underconfident, meaning that it predicts lower probabilities than the reality.
So the logits are calibrated using isotonic regression, a non-parametric method that fits a monotonic mapping from predicted probabilities to observed frequencies, fitted on the calibration set and then evaluated on the test set. This improves the match between predicted confidence and empirical accuracy. The calibrator is trained with a single pass over the calibration set; this process does not modify the learned parameters of the underlying model, it only learns the calibration mapping. The source code for calibration is in src/calibration.py and calibration.py
These are the resulting reliability curves:
Generally speaking, calibration improved reliability of each class probability. Although, there is one class, Pleural Others (I think, I'm not so sure, I'm colorblind) that has not responded well to calibration, probably due to its distribution in the test set.
Metrics computed to test the model are Precision, Recall and F1 score for each class, using non-calibrated vs. calibrated outputs.
precision recall f1-score support
No Finding 0.59 0.69 0.63 3473
Atelectasis 0.88 0.13 0.22 3928
Cardiomegaly 0.89 0.21 0.33 3372
Consolidation 0.81 0.09 0.16 2972
Edema 0.77 0.08 0.15 2577
Enlarged Cardiomediastinum 0.90 0.48 0.63 3703
Fracture 0.89 0.01 0.02 1538
Lung Lesion 1.00 0.00 0.00 1218
Lung Opacity 0.94 0.32 0.48 4857
Pleural Effusion 0.93 0.02 0.03 3547
Pleural Other 0.33 0.00 0.00 727
Pneumonia 1.00 0.01 0.01 1482
Pneumothorax 0.50 0.00 0.00 963
Support Devices 0.90 0.29 0.44 3690
micro avg 0.78 0.23 0.35 38047
macro avg 0.81 0.17 0.22 38047
weighted avg 0.85 0.23 0.30 38047
samples avg 0.40 0.30 0.32 38047
precision recall f1-score support
No Finding 0.64 0.57 0.60 3473
Atelectasis 0.72 0.67 0.70 3928
Cardiomegaly 0.72 0.73 0.72 3372
Consolidation 0.69 0.67 0.68 2972
Edema 0.63 0.66 0.64 2577
Enlarged Cardiomediastinum 0.78 0.90 0.83 3703
Fracture 0.57 0.13 0.21 1538
Lung Lesion 0.76 0.03 0.06 1218
Lung Opacity 0.81 0.72 0.76 4857
Pleural Effusion 0.64 0.60 0.62 3547
Pleural Other 0.00 0.00 0.00 727
Pneumonia 0.70 0.07 0.12 1482
Pneumothorax 0.00 0.00 0.00 963
Support Devices 0.77 0.76 0.76 3690
micro avg 0.72 0.60 0.65 38047
macro avg 0.60 0.46 0.48 38047
weighted avg 0.68 0.60 0.61 38047
samples avg 0.48 0.48 0.46 38047
Micro-average AUROC: 0.869
Macro-average AUROC: 0.842
The model shows uneven performance across the different classes. Without calibration, some conditions like Atelectasis or Cardiomegaly achieve high precision but very low recall, meaning that the model often misses true positive cases. For rarer classes such as Fracture, Lung Lesion, Pleural Other, or Pneumothorax, the model fails completely; this highlights the challenges of predicting rare conditions in a multi-label dataset and the limitations of a small trainable network.
Applying isotonic regression calibration clearly improves the model. Recall increases substantially for many classes while precision remains largely stable, leading to higher overall F1-scores. The predicted probabilities now better match actual outcomes, making the model’s predictions more reliable. Importantly, this improvement comes entirely from post-processing; the ResNet backbone remains frozen, and the model size does not increase. Calibration improved the balance and confidence of predictions without the computational cost of training a larger or more complex model.
The ROC curves and AUROC scores confirm that the model can discriminate well between positive and negative cases, with a micro-average AUROC of 0.87 and a macro-average of 0.84. Calibration mainly adjusts predicted probabilities to better reflect real-world likelihoods; it does not fundamentally change the model’s discriminative ability; rare classes remain challenging.
Given the architecture, where most parameters come from a frozen ResNet and only a small fraction is trainable, the model could benefit from unfreezing some of the final ResNet layers or experimenting with a larger backbone, such as ResNet34 or ResNet50. Improving class balance could also boost performance, especially for underrepresented diseases.
Overall, the model provides a good balance between efficiency and reliability. It produces trustworthy probability estimates and demonstrates that meaningful improvements can be achieved with careful calibration. While the average F1-score of around 0.65 may not compete with state-of-the-art benchmarks, it shows that the architecture is solid and can be further optimized with targeted improvements. The main limitations of this project are the limited trainable parameters in the model, the difficulty in predicting rare classes due to data imbalance, and the underfitting observed during training, which could be addressed with more epochs or a larger backbone.