Stabil diffusion för kontrafaktisk bildgenerering
Cardiovascular disease is the leading cause of death worldwide, hence, computer-aided research has been increasingly more popular in the field. Medical AI could be used to find early signs of disease, diagnose, and simulate disease effects in scans. With a medical image generation model, one could potentially use it for educational purposes during medical training, but also to understand how diseases affect the body and even potential causes. Image generation models are generally trained on natural images, but this project aimed to fine-tune a stable diffuser on medical chest x-rays. Counterfactuals is an image-to-image generation aiming to transition the image from one label to another. In this case, counterfactuals were generated with cardiomegaly present or not, depending on the original label. Datasets were generated from a chest X-ray dataset to represent cardiomegaly and non-cardiomegaly patients in equal numbers. Many diffusion models have the option of segmentation masks, which could minimize the change in areas irrelevant to the disease. Segmentation methods were tested, and the approach led to HybridGNet, which can generate both lung and heart segmentation masks and is therefore multipurpose depending on the application. Multiple pretrained classification models were tested and fine-tuned for a binary cardiomegaly classification. The search for the best approach for the type of task did not show a clear winner since the different ResNet, DenseNet, and EfficientNet models performed approximately equally, around 83-85\% test accuracy. The best test accuracy was for EfficientNet V2, and it was therefore used to evaluate the generated images. The computational demands and training times did, however, differ, and for practical applications, ResNet18 and EfficientNetB0 will suffice because of their low computational demands, low training time, and comparable test accuracy to the other models. For the generative models, the Stable Diffusion v1-4 was used to fine-tune. The small dataset and lower computational power did not lead to great results with an SSIM of 0.12 and a relative SSIM of around 0.29. The counterfactual images were generated with varying strengths and tested to see the class switching abilities using the classification model. The results never achieved the class switching abilities while still maintaining similarities between real and synthetic images. The testing was, however, done on relatively small datasets, and the batch size had to be kept lower to not computationally exhausting the model. Although it did not fulfill the goal of being able to perform class switching, it shows promise for further testing. More data is needed than just the dataset, as it does not contain many patients with cardiomegaly. Combined with better computing power, this could potentially lead to improved results. This could in the future be used for teaching purposes. With further development, it could include features such as patient gender, age, or any other underlying diseases.
Publication Details
- Journal
- KTH Publication Database DiVA (KTH Royal Institute of Technology)
- Published
- 2025-01-01
- Primary Topic
- Scientific Computing and Data Management
- Type
- article
- Field-Weighted Citation Impact
- 0.00