# ML-Quiz-XRay-ReportGeneration

# ML-Quiz-XRay-ReportGeneration

Finetune a Multimodal Model for X-Ray Radiology Report Generation

Finetune a Multimodal Model for X-Ray Radiology Report Generation

This repository finetunes the LLaVA (Large Language and Vision Assistant) on IU X-Ray dataset. The `XRay-ReportGeneration.pdf` defines two Tasks:

This repository finetunes the LLaVA (Large Language and Vision Assistant) on IU X-Ray dataset. The `XRay-ReportGeneration.pdf` defines two Tasks:

- Task 1. Prompt engineering: reorganize the X-Ray report findings into predefined anatomical regions

- Task 1. Prompt engineering: reorganize the X-Ray report findings into predefined anatomical regions

- Task 2. Efficient Model Fine-tuning of LLaVA on the provided training set for report generation using parameter efficient fine-tuning methods (e.g., LoRA).

- Task 2. Efficient Model Fine-tuning of LLaVA on the provided training set for report generation using parameter efficient fine-tuning methods (e.g., LoRA).

## Environments and Requirements

## Environments and Requirements

- [Google Colab](https://colab.research.google.com/)

- [Google Colab](https://colab.research.google.com/)

- A100 GPU and 100 compute units

- A100 GPU and 100 compute units

Setup and Install:

Setup and Install:

1. Mount to your google drive account

1. Mount to your google drive account

```setup

```setup

from google.colab import drive

from google.colab import drive

drive.mount('/content/drive')

drive.mount('/content/drive')

```

```

2. Load Quiz repo from Github with finetuned model checkpoints

2. Load Quiz repo from Github with finetuned model checkpoints

```setup

```setup

!git clone https://github.com/Shen16/ML-Quiz-XRay-ReportGeneration.git

!git clone https://github.com/Shen16/ML-Quiz-XRay-ReportGeneration.git

%cd ML-Quiz-XRay-ReportGeneration

%cd ML-Quiz-XRay-ReportGeneration

```

```

3. Load `LLaVa` repo and install the repository in editable mode

3. Load `LLaVa` repo and install the repository in editable mode

```setup

```setup

!git clone https://github.com/haotian-liu/LLaVA.git

!git clone https://github.com/haotian-liu/LLaVA.git

!cd LLaVA && pip install --upgrade pip && pip install -e .

!cd LLaVA && pip install --upgrade pip && pip install -e .

```

```

4. Install flash attention and DeepSpeed

4. Install flash attention and DeepSpeed

```setup

```setup

!cd LLaVA && pip install -e ".[train]"

!cd LLaVA && pip install -e ".[train]"

!pip install flash-attn --no-build-isolation

!pip install flash-attn --no-build-isolation

!pip install deepspeed

!pip install deepspeed

```

```

5. Install other packages

5. Install other packages

```setup

```setup

!pip install -r requirements.txt

!pip install -r requirements.txt

```

```

## Dataset

## Dataset

- Data was obtained from the [IU-X-ray](https://paperswithcode.com/dataset/iu-x-ray) datatset. 

- Data was obtained from the [IU-X-ray](https://paperswithcode.com/dataset/iu-x-ray) datatset. 

- Download the raw dataset from this [link](https://drive.google.com/file/d/1lBEpxrmwBkgVTZGZ92mCu0qcoq58-yMI/view?usp=sharing)

- Download the raw dataset from this [link](https://drive.google.com/file/d/1lBEpxrmwBkgVTZGZ92mCu0qcoq58-yMI/view?usp=sharing)

## Preprocessing

## Preprocessing

Description of preprocessing method:

Description of preprocessing method:

- `prompt.ipynb` notebook was used to generate the categorized reports on the validation split to complete Task 1. The output is saved in `annotation.json` file.

- `prompt.ipynb` notebook was used to generate the categorized reports on the validation split to complete Task 1. The output is saved in `annotation.json` file.

- Processed the dataset to create single image and report pairs for each patient. The `0.png` image was used only. 

- Processed the dataset to create single image and report pairs for each patient. The `0.png` image was used only. 

*Note: Did not load multiple images for a single patient for GPU and memory limitations).*

*Note: Did not load multiple images for a single patient for GPU and memory limitations).*

Running the data preprocessing code:

Running the data preprocessing code:

1. Run the script to generate the report for Task 1. Get your OpenAI API token and replace `<openAI_API_key>` with it:

1. Run the script to generate the report for Task 1. Get your OpenAI API token and replace `<openAI_API_key>` with it:

```python

```python

!python prompt.py -i 'data/annotation_quiz_all.json' -o 'data/annotation.json' -k "<openAI_API_key>"

!python prompt.py -i 'data/annotation_quiz_all.json' -o 'data/annotation.json' -k "<openAI_API_key>"

```

```

*Alternatively you can run the `prompt.ipynb` notebook in colab*

*Alternatively you can run the `prompt.ipynb` notebook in colab*

2. Generate image-report pairs for each patient and generate `dataset.json` file in the following structure for training:

2. Generate image-report pairs for each patient and generate `dataset.json` file in the following structure for training:

```json

```json

[

[

    {

    {

        "id": "360691e3-610f-4d6d-9f71-99672524eb89",

        "id": "360691e3-610f-4d6d-9f71-99672524eb89",

        "image": "360691e3-610f-4d6d-9f71-99672524eb89.png",

        "image": "360691e3-610f-4d6d-9f71-99672524eb89.png",

        "conversations": [

        "conversations": [

            {

            {

                "from": "human",

                "from": "human",

                "value": "Please describe the findings in the X-ray."

                "value": "Please describe the findings in the X-ray."

            },

            },

            {

            {

                "from": "gpt",

                "from": "gpt",

                "value": "<s_bone>Degenerative changes are present in the spine.</s_bone><s_heart>Heart size and pulmonary vascularity appear within normal limits.</s_heart><s_lung>Lungs are free of focal airspace disease. No pneumothorax or pleural effusion is seen.</s_lung><s_mediastinal></s_mediastinal><s_others>A large hiatal hernia is noted.</s_others>"

                "value": "<s_bone>Degenerative changes are present in the spine.</s_bone><s_heart>Heart size and pulmonary vascularity appear within normal limits.</s_heart><s_lung>Lungs are free of focal airspace disease. No pneumothorax or pleural effusion is seen.</s_lung><s_mediastinal></s_mediastinal><s_others>A large hiatal hernia is noted.</s_others>"

            }

            }

        ]

        ]

    },

    },

    {

    {

        "id": "511d0a14-3b13-4814-8658-baf1c47b4188",

        "id": "511d0a14-3b13-4814-8658-baf1c47b4188",

        "image": "511d0a14-3b13-4814-8658-baf1c47b4188.png",

        "image": "511d0a14-3b13-4814-8658-baf1c47b4188.png",

        "conversations": [

        "conversations": [

            {

            {

                "from": "human",

                "from": "human",

                "value": "Please describe the findings in the X-ray."

                "value": "Please describe the findings in the X-ray."

            },

            },

            {

            {

                "from": "gpt",

                "from": "gpt",

                "value": "<s_bone>Bony structures are intact.</s_bone><s_heart>Cardiac contours are within normal limits.</s_heart><s_lung>Lungs are clear.</s_lung><s_mediastinal>Mediastinal contours are within normal limits.</s_mediastinal><s_others></s_others>"

                "value": "<s_bone>Bony structures are intact.</s_bone><s_heart>Cardiac contours are within normal limits.</s_heart><s_lung>Lungs are clear.</s_lung><s_mediastinal>Mediastinal contours are within normal limits.</s_mediastinal><s_others></s_others>"

            }

            }

        ]

        ]

    },

    },

    ...

    ...

]

]

```

```

Each image is assigned a unique id. Organize the data as follows in `./ML-Quiz-XRay-ReportGeneration`:

Each image is assigned a unique id. Organize the data as follows in `./ML-Quiz-XRay-ReportGeneration`:

```

```

├── dataset_train

├── dataset_train

│   └── images

│   └── images

│       └── 00a7a45c-af6e-4f5b-b305-595620ae4deb.png

│       └── 00a7a45c-af6e-4f5b-b305-595620ae4deb.png

│       └── ...

│       └── ...

│   └── train

│   └── train

│       └── train_dataset.json

│       └── train_dataset.json

│ 

│ 

├── dataset_test

├── dataset_test

│   └── images

│   └── images

│       └── 00a7a45c-af6e-4f5b-b305-595620ae9ner.png

│       └── 00a7a45c-af6e-4f5b-b305-595620ae9ner.png

│       └── ...

│       └── ...

│   └── test

│   └── test

│       └── test_dataset.json

│       └── test_dataset.json

│ 

│ 

├── dataset_val

├── dataset_val

    └── images

    └── images

        └── 00a7a45c-af6e-4f5b-b305-595620ae0reb.png

        └── 00a7a45c-af6e-4f5b-b305-595620ae0reb.png

        └── ...

        └── ...

    └── val

    └── val

        └── val_dataset.json

        └── val_dataset.json

```

```

Run this script to generate the files for each split:

Run this script to generate the files for each split:

```python

```python

!python preprocess_data.py -s 'train'

!python preprocess_data.py -s 'train'

!python preprocess_data.py -s 'val'

!python preprocess_data.py -s 'val'

!python preprocess_data.py -s 'test'

!python preprocess_data.py -s 'test'

```

```

## Training

## Training

To train the model, run this notebook in Google Colab:

To train the model, run this notebook in Google Colab:

```train

```train

python finetune.ipynb

python finetune.ipynb

```

```

1. Set environment variable to reduce memory fragmentation issues

1. Set environment variable to reduce memory fragmentation issues

```

```

import os

import os

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

```

```

2. Run the train_mem.py script from LLaVA repo to finetune model. Used LoRa adapters for parameter efficient finetuning. Used liuhaotian/llava-v1.5-7b base model with 1 epoch (to stay within compute limit). Finetuned model weights saved to ./checkpoints/llava-v1.5-7b-task-lora folder.

2. Run the train_mem.py script from LLaVA repo to finetune model. Used LoRa adapters for parameter efficient finetuning. Used liuhaotian/llava-v1.5-7b base model with 1 epoch (to stay within compute limit). Finetuned model weights saved to ./checkpoints/llava-v1.5-7b-task-lora folder.

```

```

!deepspeed LLaVA/llava/train/train_mem.py \

!deepspeed LLaVA/llava/train/train_mem.py \

    --lora_enable True --lora_r 128 --lora_alpha 256 --mm_projector_lr 2e-5 \

    --lora_enable True --lora_r 128 --lora_alpha 256 --mm_projector_lr 2e-5 \

    --deepspeed LLaVA/scripts/zero3.json \

    --deepspeed LLaVA/scripts/zero3.json \

    --model_name_or_path liuhaotian/llava-v1.5-7b \

    --model_name_or_path liuhaotian/llava-v1.5-7b \

    --version v1 \

    --version v1 \

    --data_path ./dataset_train/train/train_dataset.json \

    --data_path ./dataset_train/train/train_dataset.json \

    --image_folder ./dataset_train/images \

    --image_folder ./dataset_train/images \

    --vision_tower openai/clip-vit-large-patch14-336 \

    --vision_tower openai/clip-vit-large-patch14-336 \

    --mm_projector_type mlp2x_gelu \

    --mm_projector_type mlp2x_gelu \

    --mm_vision_select_layer -2 \

    --mm_vision_select_layer -2 \

    --mm_use_im_start_end False \

    --mm_use_im_start_end False \

    --mm_use_im_patch_token False \

    --mm_use_im_patch_token False \

    --image_aspect_ratio pad \

    --image_aspect_ratio pad \

    --group_by_modality_length True \

    --group_by_modality_length True \

    --bf16 True \

    --bf16 True \

    --output_dir ./checkpoints/llava-v1.5-7b-task-lora \

    --output_dir ./checkpoints/llava-v1.5-7b-task-lora \

    --num_train_epochs 1 \

    --num_train_epochs 1 \

    --per_device_train_batch_size 16 \

    --per_device_train_batch_size 16 \

    --per_device_eval_batch_size 4 \

    --per_device_eval_batch_size 4 \

    --gradient_accumulation_steps 1 \

    --gradient_accumulation_steps 1 \

    --evaluation_strategy "no" \

    --evaluation_strategy "no" \

    --save_strategy "steps" \

    --save_strategy "steps" \

    --save_steps 50000 \

    --save_steps 50000 \

    --save_total_limit 1 \

    --save_total_limit 1 \

    --learning_rate 2e-4 \

    --learning_rate 2e-4 \

    --weight_decay 0. \

    --weight_decay 0. \

    --warmup_ratio 0.03 \

    --warmup_ratio 0.03 \

    --lr_scheduler_type "cosine" \

    --lr_scheduler_type "cosine" \

    --logging_steps 1 \

    --logging_steps 1 \

    --tf32 True \

    --tf32 True \

    --model_max_length 2048 \

    --model_max_length 2048 \

    --gradient_checkpointing True \

    --gradient_checkpointing True \

    --dataloader_num_workers 4 \

    --dataloader_num_workers 4 \

    --lazy_preprocess True \

    --lazy_preprocess True \

    --report_to tensorboard

    --report_to tensorboard

```

```

## Trained Models

## Trained Models

You can download trained model here:

You can download trained model here:

1. Download the model checkpoint folder from the [zip file](https://drive.gogithubZipfile.git). 

1. Download the model checkpoint folder from the [zip file](https://drive.gogithubZipfile.git). 

2. To load the model excute the following code:

2. To load the model excute the following code:

```python

```python

#merge the LoRA weights with the full base model

#merge the LoRA weights with the full base model

!python LLaVA/scripts/merge_lora_weights.py --model-path checkpoints/llava-v1.5-7b-task-lora --model-base liuhaotian/llava-v1.5-7b --save-model-path llava-ftmodel

!python LLaVA/scripts/merge_lora_weights.py --model-path checkpoints/llava-v1.5-7b-task-lora --model-base liuhaotian/llava-v1.5-7b --save-model-path llava-ftmodel

```

```

*Make sure the `llava-ftmodel` is in the `./ML-Quiz-XRay-ReportGeneration` directory.*

*Make sure the `llava-ftmodel` is in the `./ML-Quiz-XRay-ReportGeneration` directory.*

## Inference

## Inference

1. To infer the testing cases, run the following code on a Test sample image:

1. To infer the testing cases, run the following code on a Test sample image:

```python

```python

!python -m llava.serve.cli \

!python -m llava.serve.cli \

  --model-path llava-ftmodel \

  --model-path llava-ftmodel \

  --image-file "/content/drive/My Drive/ML-Quiz-XRay-ReportGeneration/dataset_test/images/cf33da4a-49f3-4dd1-8e5b-038d2637751f.png"

  --image-file "/content/drive/My Drive/ML-Quiz-XRay-ReportGeneration/dataset_test/images/cf33da4a-49f3-4dd1-8e5b-038d2637751f.png"

```

```

*Note: Replace image id "cf33da4a-49f3-4dd1-8e5b-038d2637751f.png" with exisiting id*

*Note: Replace image id "cf33da4a-49f3-4dd1-8e5b-038d2637751f.png" with exisiting id*

2. Generate predictions on Test set:

2. Generate predictions on Test set:

```python

```python

!python -m llava.eval.model_vqa_science \

!python -m llava.eval.model_vqa_science \

    --model-path llava-ftmodel \

    --model-path llava-ftmodel \

    --question-file ./dataset_test/test/test_dataset.json \

    --question-file ./dataset_test/test/test_dataset.json \

    --image-folder ./dataset_test/images \

    --image-folder ./dataset_test/images \

    --answers-file ./dataset_test/answers/llava-v1ft.5-7b.jsonl \

    --answers-file ./dataset_test/answers/llava-v1ft.5-7b.jsonl \

    --single-pred-prompt \

    --single-pred-prompt \

    --temperature 0 \

    --temperature 0 \

    --conv-mode vicuna_v1

    --conv-mode vicuna_v1

```

```

3. Generate predictions on Val set:

3. Generate predictions on Val set:

```python

```python

!python -m llava.eval.model_vqa_science \

!python -m llava.eval.model_vqa_science \

    --model-path llava-ftmodel \

    --model-path llava-ftmodel \

    --question-file ./dataset_val/val/val_dataset.json \

    --question-file ./dataset_val/val/val_dataset.json \

    --image-folder ./dataset_val/images \

    --image-folder ./dataset_val/images \

    --answers-file ./dataset_val/answers/llava-v1ft.5-7b.jsonl \

    --answers-file ./dataset_val/answers/llava-v1ft.5-7b.jsonl \

    --single-pred-prompt \

    --single-pred-prompt \

    --temperature 0 \

    --temperature 0 \

    --conv-mode vicuna_v1

    --conv-mode vicuna_v1

```

```

*Note: Was not able to evalute on validation set as I exhausted my colab compute units (requires A100 GPU)*

*Note: Was not able to evalute on validation set as I exhausted my colab compute units (requires A100 GPU)*

## Evaluation

## Evaluation

1. Clone the GREEN repo:

1. Clone the GREEN repo:

```python

```python

!git clone https://github.com/Stanford-AIMI/GREEN.git

!git clone https://github.com/Stanford-AIMI/GREEN.git

```

```

2. Install required packages. Restart kernel afterwards.:

2. Install required packages. Restart kernel afterwards.:

```python

```python

 # run and then restart kernel to use packages

 # run and then restart kernel to use packages

%cd GREEN

%cd GREEN

!pip install -e . 

!pip install -e . 

```

```

3. Import green3.py script which will save both GREEN summary and result dataframe:

3. Import green3.py script which will save both GREEN summary and result dataframe:

```python

```python

# import libraries

# import libraries

from src import green # import green.py

from src import green # import green.py

from src.green3 import compute # modified code (green3.py) to save both GREEN summary and result_df

from src.green3 import compute # modified code (green3.py) to save both GREEN summary and result_df

```

```

4. To compute the GREEN evaluation metrics on test set, run the notebook:

4. To compute the GREEN evaluation metrics on test set, run the notebook:

```python

```python

GREEN_eval.ipynb

GREEN_eval.ipynb

```

```

## Results

## Results

Our method achieves the following performance on XRay-ReportGeneration:

Our method achieves the following performance on XRay-ReportGeneration:

| Index | GREEN Score                                                    |      Bone      |     Heart      |     Lung       | Mediastinal   |

| Index | GREEN Score                                                    |      Bone      |     Heart      |     Lung       | Mediastinal   |

|-------|----------------------------------------------------------------|:--------------:|:--------------:|:--------------:|:-------------:|

|-------|----------------------------------------------------------------|:--------------:|:--------------:|:--------------:|:-------------:|

| 0     | Green average                                                  |  0.3253968254  |  0.8163265306  |  0.6737001944  | 0.5939625850  |

| 0     | Green average                                                  |  0.3253968254  |  0.8163265306  |  0.6737001944  | 0.5939625850  |

| 1     | Standard variation                                             |  0.4635559317  |  0.3657609376  |  0.3195053647  | 0.4767040097  |

| 1     | Standard variation                                             |  0.4635559317  |  0.3657609376  |  0.3195053647  | 0.4767040097  |

| 2     | (a) False report of a finding in the candidate                 |  0.7959183673  |  0.9081632653  |  0.5425170068  | 0.8469387755  |

| 2     | (a) False report of a finding in the candidate                 |  0.7959183673  |  0.9081632653  |  0.5425170068  | 0.8469387755  |

| 3     | (b) Missing a finding present in the reference                 |  0.9217687075  |  0.9098639456  |  0.8095238095  | 0.7942176871  |

| 3     | (b) Missing a finding present in the reference                 |  0.9217687075  |  0.9098639456  |  0.8095238095  | 0.7942176871  |

| 4     | (c) Misidentification of a finding's anatomic location/position|  1.0000000000  |  0.9897959184  |  0.9982993197  | 0.9829931973  |

| 4     | (c) Misidentification of a finding's anatomic location/position|  1.0000000000  |  0.9897959184  |  0.9982993197  | 0.9829931973  |

| 5     | (d) Misassessment of the severity of a finding                 |  1.0000000000  |  0.9778911565  |  0.9931972789  | 0.9948979592  |

| 5     | (d) Misassessment of the severity of a finding                 |  1.0000000000  |  0.9778911565  |  0.9931972789  | 0.9948979592  |

| 6     | (e) Mentioning a comparison that isn't in the reference        |  0.9982993197  |  0.9965986395  |  0.9880952381  | 0.9965986395  |

| 6     | (e) Mentioning a comparison that isn't in the reference        |  0.9982993197  |  0.9965986395  |  0.9880952381  | 0.9965986395  |

| 7     | (f) Omitting a comparison detailing a change from a prior study|  1.0000000000  |  1.0000000000  |  1.0000000000  | 1.0000000000  |

| 7     | (f) Omitting a comparison detailing a change from a prior study|  1.0000000000  |  1.0000000000  |  1.0000000000  | 1.0000000000  |

![GREEN Score Evaluation](GREEN_Score_Test.png)

![GREEN Score Evaluation](https://github.com/Shen16/ML-Quiz-XRay-ReportGeneration/blob/main/GREEN_Score_Test.png?raw=true)

## Acknowledgement

## Acknowledgement

> We thank the contributors of public datasets. 

 We thank the contributors of public datasets.