This article presents a methodology for chest CT scan analysis that enables the automatic categorization of pulmonary tuberculosis cases into one of the following five types: Infiltrative, Focal, Tuberculoma, Miliary, Fibro-cavernous. We showcase several deep learning methods for classifying tuberculosis in CT scans from the ImageCLEF 2021 Tuberculosis - TBT classification challenge. Furthermore, it explores the use of pre-trained models, as well as training from scratch on volumetric data and 2D projections.
In previous editions, diferent approaches were proposed for the tuberculosis task. In 2017, there were two important tasks: MDR Detection which consisted of assigning each TB patient the probability of having a resistant form of tuberculosis based on the analysis of chest CT scan; Detection of TB Type task which automatically categorize each TB case into one of the following five types: Infiltrative, Focal, Tuberculoma, Miliary, Fibro-cavernous.
In 2018, the valuable result was of team UIIP-BioMed that used a Deep CNN [
In 2019, UIIP-BioMed [
In 2020, the majority of the participants used some variations of the projection-based approach and created 2D CNNs. As a result, only four groups tried 3D CNNs for a direct analysis of the CT volumetric data [11]. SenticLab.UAIC [12] was the winner of the task using 2D and 3D CNNs. The SDVA-UCSD was ranked on second place using a 3D CNN with a convolutional block attention module (CBMA) and a customized loss function [13]. Our team was ranked on 7th place, using SVMs and CNNs for lung-wise processing [14].
We propose an extension to the 2D projection methods which were successfully applied in previous years by replacing the arithmetic mean aggregation function with k-means. We made this choice since it can be considered that the average is a specific case for k-means when k=1. With this, we hope to capture more relevant characteristics from the data. This gives us the advantage of working with smaller dimensionality data.
Transforming 3D CT scans into 2D images does have the downside of losing spacial proximity information along the third dimension and introduces the probability of losing important features. Due to this, we also experimented working with the 3D samples directly in order to have a reference point of how the two modalities (2D and 3D) performed.
Based on the metadata file presented by the organizers, we concluded that the data at hand is highly imbalanced. Since this could have negatively influenced our models, we have decided to use weighted loss, in order to counteract this phenomenon.
Table 1 provides an insight for the distribution of each afection in the training set, as well as the weight used for the loss function. To calculate the weight, we have divided the lowest total number of patients afected by one of the types (in our case, this would be 70) by the number of patients afected for each type.
Given that the masks provided were not able to return the best segmentation possible, a decision was made to combine the two of them and compute a new mask that would be closer to what we needed. However, this happened to be quite hard, as some of the masks did not have a proper lung segmentation, in accordance to the anatomical position given by the midsagittal line in the upper thoracic area. Furthermore, some issues were not related to the segmentations, but by the CT Scans, since some of them were not able to be opened or were considered faulty. Table 2 presents what issues we have encountered when processing the data and the proposed solution for said problems.
The most dificult case to treat was the test CT Scan for patient TST_247, because we were not able to load it properly and it needed to be processed and labeled. Since it had a similar issue to the scan of patient TRN_360, we have assumed that both files have a similar origin and that the problem was either in an underlying anatomical condition or in the acquisition of a similar device fault. Thus, we manually assigned the same TBT class that TRN_360 was assigned in the training metadata.
When normalizing the images, we have opted for a similar approach to past high-ranking
participants[
The two types of lung segmentations provided by the organizers had diferent issues. One of them has accurate edges, but it tends to overlook or miss entirely some of the large cavities or other lesions in the lungs that would be relevant for the model. However, the other performs better in covering the entire lung, but it is generally more inaccurate, especially when referring to the edges of the lung lobes. The solution proposed for this issue was to generate a mean of the two masks. Combining them and retrieving the entire information, helped ensure that no particularity of the afections is overlooked, while also highlighting the increased importance of the sections where the segmentation methods overlap and hence agree upon. To complete this task, we looked at the first type of mask, the fully automatic multistage one, taking into account only the non-zero values, added them to the second type of mask and then computed the mean of the two, multiplying it by 255 in the end. Therefore, we computed a new image that had the opacity adjusted for the parts of the lungs where the two segmentations did not overlap.
The resulting masks were applied to the raw CT scans by multiplying the two tensors. However, when completing this step, there was a threshold set so that it could help with eliminating the non-lung pixels present in the resulting image.
Multiple types of data augmentation techniques were attempted to be used with a view to prevent an early over-fitting of the models. To this end, we have employed Randaugment [15], only disabling the "Invert" and "Solarize" augmentations and setting = 2 and = 14. Additionally, we used Random Erasing [16] which is a technique that randomly crops out a patch of size between 5% and 15% from the original image. Nonetheless, there were other approaches to augmenting the dataset that we have tried to apply, such as Mixup [17] and Cutmix [18], but these methods proved to produce images which were making it too dificult for the model to learn on.
1the mask referred to is the fully automatic multistage one
When working with the 3D volumetric data, we have looked for a model that had been pretrained, preferably on medical data. Thus, we came across MedicalNet [19], which provided some interesting models that had been trained previously. We have worked with these in order to obtain better results that could not be achieved when training the network from scratch. However, due to the large dimension of the samples and the hardware limitation imposed by using a single Nvidia RTX 2080Ti GPU, we had to resize all the images to 256x256x32. Taking these facts into consideration, we managed to fit into the GPU memory only the MedicalNet10 variant, which is essentially a pre-trained 3D-ResNet10 network.
One technique adopted when training was gradient accumulation. Therefore, since we could only use a maximum batch size of 2, we accumulated 16 gradient backpropagation steps before updating the weights, efectively allowing us to simulate a batch size of 32.
We loaded the 23-dataset pre-trained MedicalNet10 weights, replacing its segmentation head with a newly initialized classification head, where a linear layer outputs the probabilities for each of the 5 classes.
As for the hyper parameter tuning step, we attempted multiple variations of optimizers and data augmentations, but unfortunately our attempts were consistently faced with an early overfitting of the model. This means that it would happen after the first 40 epochs of training, when very little to no data augmentations were applied. One other issue was that frequently we would encounter an unstable training process, caused by exploding gradients, which would then be followed by a model collapse, usually occurring in the first 5 epochs when stronger augmentations were used.
Our best submitted run is a MedicalNet10 model trained for 44 epochs using the AdamP optimizer [20] with the default parameters, cross entropy loss and with no data augmentations. Attempts to train 3D-ResNet variants from scratch also sufered from the same issues but the overfitting would start sooner and as such the models generated poorer results.
In previous years, the computation and usage of 2D projections has proven to be one of the
most successful ways of pre-processing a 3D volume. This happened mostly due to its efective
reduction in dimensionality without the loss of relevant information, so, this year, we attempted
to build upon this approach. We had managed to adapt the projection method to use the k-means
algorithm [21][22]. As described in [
Using k-means as an aggregation function gives us more flexibility with the number of potentially relevant features that we extract as centroids from the 3D volume. Since the function outputs a number of centroids equal to the value of k (not just 1, as in the case of average()), each of these centroids is mapped to a diferent channel in the resulting projection. The centroids are initialized randomly, this factor did not seem to be a big influence on the results since the algorithm always and very rapidly finds the final centroids, in very few steps, as it is a very simple 1D k-means calculation. The mapping is done by sorting the centroid values, the smallest valued centroids will be mapped to the first channel, the second smallest to the second channel and so on. Our final projection algorithm still used max() and std() as aggregation functions and replaced average() with k-means() where k=5. The resulting images consist of 7 channel projections for each of the axis directions X, Y and Z. The usage of max() could be replaced by using higher values for k since the centroid on the higher end of the interval will probably be very close in value to the max value.
An ablation study would be needed to identify an optimal value for k, as a starting point, we arbitrarily picked k=5. Since a stable training procedure was not obtained, we did not get to ifne-tune this parameter.
The Y projection had to be treated separately, since projecting from that direction would overlap the two lungs. In order to prevent losing relevant information at the lung level, we ifrst separated the lungs with the help of the lung-wise segmentation and then calculated the projection separately for each one. Consequently, we concatenated the two images such that the two lungs would be positioned next to each other. The resulting X, Y and Z projections can be observed in Fig.1. Due to having more than 3 channels, the figures here were obtained by resizing the number of channels to 3, for RGB. However, there exists less information, because some was lost in the process of resizing the channels for visualization purposes.
After computing the X, Y and Z projections, each with 7 channels, we have decided to stack them along the channel dimension. Given that the anatomical positioning is the same, overlapping the lobes would not be considered an issue, especially since there would be 21 individual channels for the model to look at and learn the particularities of each tuberculosis type. An approximated look of the resulting images is presented in Fig.2. This was favorable especially because the algorithm would only use one image for the training and testing process, meaning that the data would be fed all at once. A fixed random seed was used for to generate the k-means projections. 3.4.2. Models The considerably smaller 2D image size allowed us to use larger models and batch sizes. To do so, we tried resizing the images, while preserving the number of channels. As for the CNN models, we experimented with PreResNet56 [23] and the Swin Transformer [24] large, base and tiny variants.
For the Swin Transformer models, we also tried to use the weights trained on ImageNet21k [25] and to apply various transfer learning techniques in order to adapt it to the TBT classification task. One of the steps implemented to do this was to reinitialize the PatchEmbed layer and the final normalization, pooling and linear layers, so that they can match our 5 classes. The Swin Transformer model runs were not able to noticeably reduce the loss when trained from scratch or when loading the pre-trained weights, but we believe this might have been due to an implementation error on our part or from a faulty normalization technique.
The two models that delivered the best results were both based on PreResNet, with a depth of 56. Most of the experiments with this type of model were based on hyper-parameter tuning, as the base was a standard PreResNet56 that had been trained from scratch on data for the task.
For the hyper-parameters, we used the RAdam optimizer[26] with the default parameters, a cross entropy loss function and a batch size of 16 with 4 gradient accumulation steps, resulting in a simulated batch size of 64. Both models use lookahead[27]. Further information about the diference between the parameters of the two models is presented in Table.3.
When evaluating our experiments, the main reference was the kappa score, defined on the
interval [
The experiments were performed using only flips and rotations as data augmentations, this caused the model to overfit very early on and would not recover with more training. This phenomenon convinced us to apply stronger augmentation techniques, such as Randaugment and Random Erasing, in order to populate the latent space. However, when these stronger augmentations were being used, the models sufered from training instability since our loss values would explode and hence the gradients would too, resulting in numeric overflow. Finally, once the loss exploded, the model collapsed and did not learn anything else. This happened very fast, usually in the first 5 epochs of training.
The four submissions made by our team are listed in Table.4, detailing what model was used, whether or not the model had pre-trained weights, the local validation kappa score and the test kappa score. Since training with strong data augmentations would result in an unstable training, which stopped the model from learning in the first few epochs, no submissions were made with them. All of the runs we submitted sufered from relatively fast overfitting, this occurring in the first approximately 50 epochs.
Further work needs to be invested into adapting state-of-the-art data augmentations from the general computer vision domain to medical images. Since the size of the datasets usually encountered in the field is small, stable data augmentation strategies are required, as frequently acquiring additional data is not feasible.
As a retrospective analysis of our methods, we need to carefully investigate the exact causes of training instability. We suspect that, because we did not normalize the data after the final step in the pre-processing/augmentation pipeline, the resulting images might have very diferent value distributions between diferent samples, which could even escape our intended value interval of [0.0, 1.0]. Another potential cause of our issues could come from the large size of the linear classification layer which could be the cause of the exploding gradients we noticed, we have to investigate if using a pooling layer before the classification one helps solve our issue.
In this paper we present ways of detecting diferent types of tuberculosis in lungs. The methods that have proven to work best for us are focused on pre-trained models for classifying 3D volumetric data and models trained from scratch for 2D projections.
The best result from our submissions can be improved, but in doing so, we would have to investigate diferent normalization techniques and types of data augmentation that would suit the task better. However, in order to work with volumetric data properly, we would also need to improve the current hardware with the aim of being able to apply larger models, as well as greater batch sizes, on medical datasets.
This work was supported by project REVERT (taRgeted thErapy for adVanced colorEctal canceR paTients), Grant Agreement number: 848098, H2020-SC1-BHC-2018-2020/H2020-SC1-2019Two-Stage-RTD. [11] S. Kozlovski, V. Liauchuk, Y. Dicente Cid, A. Tarasau, V. Kovalev, H. Muller, Overview of imagecleftuberculosis 2020 - automatic ct–based report generation, CLEF2020 Working Notes (2020). [12] R. Miron, C. Moisii, M. Breaban, Revealing lung afections from cts. a comparative analysis of various deep learning approaches for dealing with volumetric data, CLEF2020 Working Notes (2020). [13] X. Lu, E. Chang, Z. Liu, C. Hsu, J. Du, A. Gentili, ImageCLEF2020: Laterality-reduction three-dimensional cbam-resnet with balanced sampler for multi-binary classification of tuberculosis and ct auto reports, CLEF2020 Working Notes (2020). [14] L. Coca, A. Hanganu, C. Cusmuliuc, A. Iftene, UAIC2020: Lung analysis for tuberculosis detection, CLEF2020 Working Notes (2020). [15] E. D. Cubuk, B. Zoph, J. Shlens, Q. V. Le, Randaugment: Practical automated data augmentation with a reduced search space, in: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops, 2020, pp. 702–703. [16] Z. Zhong, L. Zheng, G. Kang, S. Li, Y. Yang, Random erasing data augmentation, in: Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, 2020, pp. 13001– 13008. [17] H. Zhang, M. Cisse, Y. N. Dauphin, D. Lopez-Paz, mixup: Beyond empirical risk minimization, arXiv preprint arXiv:1710.09412 (2017). [18] S. Yun, D. Han, S. J. Oh, S. Chun, J. Choe, Y. Yoo, Cutmix: Regularization strategy to train strong classifiers with localizable features, in: Proceedings of the IEEE/CVF International Conference on Computer Vision, 2019, pp. 6023–6032. [19] S. Chen, K. Ma, Y. Zheng, Med3D: Transfer learning for 3D Medical Image Analysis, arXiv preprint arXiv:1904.00625 (2019). [20] B. Heo, S. Chun, S. J. Oh, D. Han, S. Yun, G. Kim, Y. Uh, J.-W. Ha, Adamp: Slowing down the slowdown for momentum optimizers on scale-invariant weights, in: Proceedings of the International Conference on Learning Representations (ICLR), Online, 2021, pp. 3–7. [21] S. P. Lloyd, Least squares quantization in pcm, IEEE Trans. Inf. Theory 28 (1982) 129–136. [22] J. MacQueen, et al., Some methods for classification and analysis of multivariate observations, in: Proceedings of the fifth Berkeley symposium on mathematical statistics and probability, volume 1, Oakland, CA, USA, 1967, pp. 281–297. [23] K. He, X. Zhang, S. Ren, J. Sun, Deep residual learning for image recognition, in: Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 770–778. [24] Z. Liu, Y. Lin, Y. Cao, H. Hu, Y. Wei, Z. Zhang, S. Lin, B. Guo, Swin transformer: Hierarchical vision transformer using shifted windows, arXiv preprint arXiv:2103.14030 (2021). [25] J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, L. Fei-Fei, Imagenet: A large-scale hierarchical image database, in: 2009 IEEE conference on computer vision and pattern recognition, Ieee, 2009, pp. 248–255. [26] L. Liu, H. Jiang, P. He, W. Chen, X. Liu, J. Gao, J. Han, On the variance of the adaptive learning rate and beyond, CoRR abs/1908.03265 (2019). URL: http://arxiv.org/abs/1908. 03265. arXiv:1908.03265. [27] M. R. Zhang, J. Lucas, G. E. Hinton, J. Ba, Lookahead optimizer: k steps forward, 1 step back, CoRR abs/1907.08610 (2019). URL: http://arxiv.org/abs/1907.08610. arXiv:1907.08610.