Comparison of Tree-Based Machine Learning Models for Classification of Tuberculosis Outcomes in Brazil
Keywords:
Tuberculosis, machine learning, clinical outcomesAbstract
Tuberculosis remains a significant public health concern, recognized as a reemerging disease strongly associated with socioeconomic conditions. According to the World Health Organization, tuberculosis continues to be the leading cause of death from a single infectious agent worldwide in 2025. This study evaluates Random Forest, XGBoost, CatBoost, and LightGBM for classifying four treatment outcomes (cure, abandonment, death from tuberculosis, and death from other causes) using 53,656 epidemiological records from Minas Gerais, Brazil (2010–2024). A preprocessing pipeline was designed to handle heterogeneous and high-cardinality variables. Models were assessed with and without SMOTE balancing under 5-fold stratified cross-validation, using accuracy, weighted F1-score, and per-class recall as evaluation metrics. Without balancing, all models achieved approximately 0.75 accuracy but failed to detect minority outcomes. Experiments across multiple random seeds, feature subset analysis, and temporal validation were conducted to assess model robustness. With SMOTE, CatBoost achieved the best overall performance, with the highest cross-validation F1-score (0.7071 ± 0.0034) and improved recall for abandonment and death outcomes. Results indicate that the combination of
clinical and socioeconomic features is essential for predictive performance, and that data quality and class imbalance are the main obstacles to reliable minority-class detection in tuberculosis surveillance.
Downloads
References
C. Busatto, D. V. Bierhals, J. S. Vianna, P. E. A. Silva, L. G. Possuelo, and I. B. Ramis, “Epidemiology and control strategies for tuberculosis in countries with the largest prison populations,” Rev. Soc. Bras. Med. Trop., vol. 55, e0060-2022, 2022, doi: https://doi.org/10.1590/0037-8682-0060-2022
World Health Organization, Global Tuberculosis Report 2025. Geneva, Switzerland: World Health Organization, 2025. [Online]. Available: https://www.who.int/publications/i/item/9789240116924
J. R. Queiroz et al., “Tendˆencia da mortalidade por tuberculose e relação com o índice sociodemogr´afico no Brasil entre 2005–2019,” Cien. Saude Colet., vol. 29, e00532023, 2024, doi: https://doi.org/10.1590/141381232024295.00532023
J. Chakaya, M. Khan, F. Ntoumi, E. Aklillu, R. Fatima, P. Mwaba, N. Kapata, S. Mfinanga, S. E. Hasnain, P. D. M. C. Katoto, A. N. H. Bulabula, N. A. Sam-Agudu, J. B. Nachega, S. Tiberi, T. D. McHugh, I. Abubakar, and A. Zumla, “Global Tuberculosis Report 2020—Reflections on the global TB burden, treatment and prevention efforts,” Int. J. Infect. Dis., vol. 113, suppl. 1, pp. S7–S12, 2021, doi: https://doi.org/10.1016/j.ijid.2021.02.107
B. Biggs, L. King, S. Basu, and D. Stuckler, “Is wealthier always healthier? The impact of national income level, inequality, and poverty on public health in Latin America,” Soc. Sci. Med., vol. 71, no. 2, pp. 266–273, 2010, doi: https://doi.org/10.1016/j.socscimed.2010.04.002
G. Freitas et al., “Completude dos dados de notificac¸˜ao e de acompanhamento da tuberculose pulmonar em Belo Horizonte, Minas Gerais, de 2001 a 2020: estudo transversal,” REME Rev. Min. Enferm., vol. 28, p.1563, 2025, doi: 10.35699/2316-9389.2025.49258
S. Bengesi, H. El-Sayed, M. K. Sarker, Y. Houkpati, J. Irungu, and T. Oladunni, “Advancements in generative AI: A comprehensive review of GANs, GPT, autoencoders, diffusion models, and transformers,” IEEE Access, vol. 12, pp. 69812–69837, 2024, doi: https://doi.org/10.1109/ACCESS.2024.3397775
L. Schut, N. Tomaˇ sev, T. McGrath, D. Hassabis, U. Paquet, and B. Kim, “Bridging the human–AI knowledge gap through concept discovery and transfer in AlphaZero,” Proc. Natl. Acad. Sci. U.S.A., vol. 122, no. 13, e2406675122, 2025, doi: https://doi.org/10.1073/pnas.2406675122
M. Shin, J. Kim, B. van Opheusden, and T. L. Griffiths, “Superhuman artificial intelligence can improve human decision-making by increasing novelty,” Proc. Natl. Acad. Sci. U.S.A., vol. 120, no. 12, e2214840120, 2023, doi: https://doi.org/10.1073/pnas.2214840120
E. Sanches, F. Augusto, P. Silveira, B. Feij´o Junqueira, D. A. M. Lemes, J. Picolo, G. Ribeiro Sales, V. Corso, and C. S. Bezerra, “Artificial intelligence for gas leak detection with thermal cameras and metal oxide semiconductor sensors,” Rev. Inform. Teor. Apl., vol. 32, no. 1, pp. 91–98, 2025, doi: https://doi.org/10.22456/2175-2745.143526
F. Kitsios, M. Kamariotou, A. I. Syngelakis, and M. A. Talias, “Recent advances of artificial intelligence in healthcare: A systematic literature review,” Appl. Sci., vol. 13, no. 13, 7479, 2023, doi: https://doi.org/10.3390/app13137479
A. Ahmadi and N. Rabie Nezhad Ganji, “AI-driven medical innovations: Transforming healthcare through data intelligence,” Int. J. BioLife Sci., vol. 2, no. 2, pp. 132–142, 2023, doi: https://doi.org/10.22034/ijbls.2023.185475
Z. Strika, K. Petkovic, R. Likic, and R. Batenburg, “Bridging healthcare gaps: A scoping review on the role of artificial intelligence, deep learning, and large language models in alleviating problems in medical deserts,” Postgrad. Med. J., vol. 101, no. 1191, pp. 4–16, 2024, doi: https://doi.org/10.1093/postmj/qgae122
L. Marques, B. Costa, M. Pereira, A. Silva, J. Santos, L. Saldanha, I. Silva, P. Magalh˜aes, S. Schmidt, and N. Vale, “Advancing precision medicine: A review of innovative in silico approaches for drug development, clinical pharmacology and personalized healthcare,” Pharmaceutics, vol. 16, no. 3, 332, 2024, doi: https://doi.org/10.3390/pharmaceutics16030332
H. Taherdoost and A. Ghofrani, “AI’s role in revolutionizing personalized medicine by reshaping pharmacogenomics and drug therapy,” Intell. Pharmacy, vol. 2, no. 5, pp. 643–650, 2024, doi: https://doi.org/10.1016/j.ipha.2024.08.005
T. A. Vicentini and L. C. Bertucci Barbosa, “Machine learning applied to biological sequence comparison: An alignment-free approach,” Rev. Inform. Teor. Apl., vol. 32, no. 2, pp. 22–35, 2025, doi: https://doi.org/10.22456/2175-2745.141585
M. Bonal, C. Padilla, G. Chevillard, and V. Lucas-Gabrielli, “A French classification to describe medical deserts: A multi-professional approach based on the first contact with the healthcare system,” Int. J. Health Geogr., vol. 23, no. 1, 5, 2024, doi: https://doi.org/10.1186/s12942-02400366-7
S. Hansun, A. Argha, I. Bakhshayeshi, A. Wicaksana, H. Alinejad Rokny, G. J. Fox, S.-T. Liaw, B. G. Celler, and G. B. Marks, “Diagnostic performance of artificial intelligence-based methods for tuberculosis detection: Systematic review,” J. Med. Internet Res., vol. 27, e69068, 2025, doi: https://doi.org/10.2196/69068
C. Cintron, M. Dauphinais, X. Du, A. Tabackman, A. Lenart, A. Laliberte, J. Dirksen, and P. Sinha, “Enriching tuberculosis research by measuring poverty better: A perspective,” BMC Glob. Public Health, vol. 3, p. 17, 2025, doi: https://doi.org/10.1186/s44263-025-00127-z
Minist´erio da Sa´ude do Brasil, Sistema de Informac¸˜ao de Agravos de Notificac¸˜ao (SINAN). Bras´ ılia, Brazil: Minist´erio da Sa´ude, 2024. [Online]. Available: https://portalsinan.saude.gov.br/ [Accessed: Jan. 08, 2026].
Brasil, Minist´erio da Sa´ude, DATASUS: Departamento de Inform´atica do Sistema ´Unico de Sa´ude. Brasília, Brazil: Minist´erio da Sa´ude, 2025. [Online]. Available: https://datasus.saude.gov.br/ [Accessed: Jan. 08, 2026].
Elsevier, Scopus: Abstract and Citation Database. Amsterdam, Netherlands: Elsevier, 2024. [Online]. Available: https://www.scopus.com [Accessed: Jan. 08, 2026].
U.S. National Library of Medicine, Montgomery County X-ray Set. Bethesda, MD, USA: National Library of Medicine, 2023. [Online]. Available: https://pubmed.ncbi.nlm.nih.gov/ [Accessed: Jan. 08, 2026].
Association for Computing Machinery, ACM Digital Library. New York, NY, USA: ACM, 2024. [Online]. Available: https://dl.acm.org [Accessed: Jan. 08, 2026].
Minas Gerais, Governo do Estado; Fundac¸˜ao de Amparo `a Pesquisa do Estado de Minas Gerais; and Universidade Federal de Minas Gerais, Centro de Desenvolvimento e Planejamento Regional, DataViva: Plataforma de visualizac¸˜ao de dados socioeconˆomicos. Belo Horizonte, Brazil, 2025. [Online]. Available: https://dataviva.info/ [Accessed: Jan. 08, 2026].
H. de Almeida Pereira, M. Roberto Ribeiro, and C. Aparecido Leite Nametala, “Comparison of Tree-Based Machine Learning Models for Classification of Tuberculosis Outcomes in Brazil,” Zenodo, v1.0.0, 2025, doi: https://doi.org/10.5281/zenodo.20012729
L. Breiman, “Random forests,” Mach. Learn., vol. 45, no. 1, pp. 5–32, 2001, doi: https://doi.org/10.1023/A:1010933404324
T. Chen and C. Guestrin, “XGBoost: A scalable tree boosting system,” in Proc. 22nd ACM SIGKDD Int. Conf. Knowl. Discov. Data Min., 2016, pp. 785–794, doi: https://doi.org/10.1145/2939672.2939785
A. V. Dorogush, V. Ershov, and A. Gulin, “CatBoost: Gradient boosting with categorical features support,” arXiv:1810.11363, 2018. doi: https://doi.org/10.48550/arXiv.1810.11363
G. Ke, Q. Meng, T. Finley, T. Wang, W. Chen, W. Ma, Q. Ye, and T.-Y. Liu, “LightGBM: A highly efficient gradient boosting decision tree,” in Proc. 31st Int. Conf. Neural Inf. Process. Syst., 2017, pp. 3149–3157.
G. Van Rossum and F. L. Drake Jr., Python Reference Manual, Python Software Foundation, 2009. Available: https://www.python.org.
T. Kluyver et al., “Jupyter Notebooks– a publishing format for reproducible computational workflows,” in Positioning and Power in Academic Publishing: Players, Agents and Agendas, IOS Press, 2016, pp. 87–90. doi: https://doi.org/10.3233/978-1-61499-649-1-87
W. McKinney, “Data Structures for Statistical Computing in Python,” in Proc. 9th Python in Science Conf. (SciPy), 2010, pp. 56–61. doi: https://doi.org/10.25080/Majora-92bf1922-00a
C. R. Harris et al., “Array programming with NumPy,” Nature, vol. 585, pp. 357–362, 2020. doi: https://doi.org/10.1038/s41586-020-2649-2
T. Solc, Unidecode: ASCII Transliterations of Unicode Text, 2024. Available: https://github.com/avian2/unidecode
E. Gazoni and C. Clark, openpyxl: A Python Library to Read/Write Excel 2010 xlsx/xlsm Files, 2024. Available: https://openpyxl.readthedocs.io
F. Pedregosa et al., “Scikit-learn: Machine Learning in Python,” J. Mach.Learn. Res., vol. 12, pp. 2825–2830, 2011.
P. Virtanen et al., “SciPy 1.0: Fundamental algorithms for scientific computing in Python,” Nature Methods, vol. 17, pp. 261–272, 2020, doi: https://doi.org/10.1038/s41592-019-0686-2
G. Lemaˆ ıtre, F. Nogueira, and C. K. Aridas, “Imbalanced-learn: A Python Toolbox to Tackle the Curated List of Selection Mechanisms for Data Sampling,” J. Mach. Learn. Res., vol. 18, no. 17, pp. 1–5, 2017. Available: http://jmlr.org/papers/v18/16-365.html
J. D. Hunter, “Matplotlib: A 2D graphics environment,” Comput. Sci. Eng., vol. 9, no. 3, pp. 90–95, 2007. doi: https://doi.org/10.1109/MCSE.2007.55
M. Waskom, “Seaborn: Statistical data visualization,” J. Open Source Softw., vol. 6, no. 60, p. 3021, 2021, doi: https://doi.org/10.21105/joss.03021
R. Story et al., folium: Python Data, Leaflet.js Maps, 2013. Available: https://github.com/python-visualization/folium
T. A. A. Pontes, F. Fernandez-Llimos, and A. Wiens, “Impact of COVID 19 on tuberculosis notifications,” Rev. Inst. Med. Trop. Sao Paulo, vol. 66, e37, 2024, doi: https://doi.org/10.1590/S1678-9946202466037
M. H. Chitwood, N. A. Menzies, P. Bartholomay, D. M. Pelissari, J. N. B. Silva Junior, L. O. Harada, F. D. C. Johansen, E. L. N. Maciel, M. C. Castro, M. Sanchez, J. L. Warren, and T. Cohen, “Quantifying disruptions to tuberculosis case detection in Brazilian states during the COVID-19 pandemic,” Int. J. Epidemiol., vol. 54, no. 5, dyaf146, 2025, doi: https://doi.org/10.1093/ije/dyaf146