BurnSight AI is a chatbot that integrates a CNN with an LLM to classify burn severity and generate first-aid recommendations respectively. See a demo here!
My motivation for working on this project started when I noticed Gemini struggling to confidently identify common injuries based on the image alone. To close this gap, I wanted to integrate an image classification model with a Gemini API such that Gemini is just given textual input of a score consisting of the classified injury and the confidence level (probability of it being that injury) where based on the inputted score, Gemini would output a recommended care plan if there was a high enough confidence score (set to be 67% though it may later change). This is to ensure that when the model is uncertain, it does not give a potentially incorrect classification.
Originally, I wanted to build a more generalized "First-Aid" chatbot that given an image of an injury, it would classify the injury and give "First-Aid" type advice on how to treat it before recieving medical attention. The user could also ask additional questions if necessary.
However, an issue I faced while developing the classifier was data. There was a lack of datasets of common First-Aid injuries such as bruises and burns on sites such as RoboFlow and Kaggle. I alternatively could've created my own dataset but as the focus of the project was medical injuries, with poor data, the model could risk in classifying an injury incorrectly leading to wrong treatment and therefore, worsening the injury.
Due to these issues, I decided to narrow down my focus to only burn injuries and classifying the severity of each burn. From all the injuries I could've focused on, I chose burns because burns are the fourth most common injury after road traffic accidents, falls, and physical violence [1]. Additionally, when it comes to treating burns, First-Aid plays a critical role by reducing burn complications such as lowering the risk of infection. [2].
-
Mobile App for Burn Classification: Similar project that uses YOLO instead of a CNN for burn severity classification. Does not integrate an LLM for treatment recommendation.
-
AI-Driven Integrated System for Burn Depth Prediction With Electronic Medical Records: Algorithm Development and Validation: Uses a VLM for burn classification. Requires ultrasound images and expert instruction to properly classify the burn. While highly suitable for clinical settings, it relies on information that is generally unavailable during first-aid situations, making it less applicable for consumer use.
-
AI-based burn image assessment: Reliability and clinical error patterns of multimodal large language models in a repeated-inference study: Compares the performance of multiple multimodal models on burn severity. Gemini reported
$76.4 \pm 6.8$ percent accuracy on the dataset. This was the highest standard deviation value from all the models tested showing that Gemini alone has high variability for burn classification.
- App: Streamlit
- Machine Learning: PyTorch, Transformers (Hugging Face), Torchvision, Pillow (PIL):
- LLM: Google Gemini API
- Session Initialization: When the application loads, a new session id (uuid) is created for the current session of the application. The session id remains the same if the application is reloaded. This is because each session id is unique to the application session. A new json file storing history is created for each session id.
- Image Upload: The user uploads an image of a burn the frontend. The image is only stored for classification purposes and deleted after a classification is made.
- Image Preprocessing: The uploaded image is passed directly to the Burn Classifier model (fine-tuned MOBILENET-V3). The image is resized and normalized according to the preprocessing requirements of the model.
-
Burn Classification: The pretrained model uses the image input to output
$3$ logit values. These$3$ logit values are passed through a softmax to find the probabilities for each of the burn severities. Then the maximum probability is decided as the classification. - LLM Processing: The Gemini API is prompted as a medical assistant with inputs being the confidence and prediction of the pretrained model. These inputs are used to generate recommended First-Aid methods for each burn.
- Conversation Initialization: After the initial diagnosis, a unique conversation id is created (uuid)(essentially anytime you open a new chat). After each chat it verifies if it belongs to an existing conversation, else it creates a new one. Each time the chat is added to json file for the respective session id history file.
- Follow-Up Questions: The user can ask follow-up questions about the diagnosed burn. These messages are associated with the current conversation ID, allowing the application to maintain the context of the original diagnosis when generating subsequent responses. Google search was initally integrated for verification of results but later removed as it enabled the LLM to answer off-topic questions. A future fix is TBD.
- History Deletion: When running, if a history json file is older than a day.
Run the following code to load the dataset used for training our image classification model:
!pip install roboflow
from roboflow import Roboflow
rf = Roboflow(api_key="API KEY")
project = rf.workspace("first-aid-vfyay").project("first-aid-app-2fa7k")
version = project.version(4)
dataset = version.download("folder")
For more details visit the RoboFlow here. You will need to access RoboFlow to get the API key to access the dataset.
A lot of this information is also available on RoboFlow but here is a brief overview.
The dataset consists of
Training: {'First degree burn': 0, 'Third degree burn': 1, 'second degree burn': 2}
Validation: {'First degree burn': 0, 'Third degree burn': 1, 'second degree burn': 2}
Testing: {'First degree burn': 0, 'Third degree burn': 1, 'second degree burn': 2}
So, if the model predicts the value
The graph above shows the distribution of data across the different classes. As we can see it is relatively uniform suggesting an equal distribution between the different burn severities. This suggests that accuracy is a good metric to evaluate the performance of the model.
RoboFlow had provided an already split dataset and claims to follow an
Training: 2595
Validation: 323
Testing: 325
Mathematically, we can verify that RoboFlow is correct.
Training:
Validation:
Testing:
The images are of dimensions
Though recent research has shown strong performance with vision transformers for many computer vision tasks, they are data hungry. Our dataset only consists of a bit over 3000 images and will cause the vision transformer to overfit. Therefore, I chose to go with a simpler approach and work with Convolutional Neural Networks (CNNs). Unlike CNNs, vision transformers also don't have a built-in understanding of local structures (i.e. Inductive Bias). This is essential because when classifying burns, local characteristics can potentially reveal patterns important for classifying particular burns such as blistering.
Instead of creating a CNN from scratch, we wanted to use transfer learning. Transfer Learning is essentially when you use a pretrained model on a large dataset and then train only a smaller portion of the model on a smaller dataset to get a strong performance for specific tasks. Aside from the learning experience, we chose transfer learning for the lower training time and stronger performance than if we trained a CNN from scratch.
For this project we evaluated two pretrained models, both of which were trained on the IMAGENET dataset.
- ResNet-50: Standard benchmark model for image classification. Consists of 25.6 million parameters.
- MOBILENET-V3: A more lightweight CNN built for mobiles and edge devices (e.g. phones and microcontrollers). Known for fast inference times but possiblilty for lower performance than ResNet-50. Has only 3 million parameters.
To see what hyperparameters were fine-tuned and what layers we froze, please visit the linked notebooks at the top of this section.
Note: To learn more about the Transformers library, we chose to use that for transfer learning of ResNet-50 instead of the torchvision counterpart. Both models have the same architecture so ideally there should be no difference in performance.
One of the most important deciding factors for model selection in our scenario would be recall. Recall that "recall" (pun not intented) is calculated by the following formula:
Below is a classification report comparing both models on the test data with multiple metrics including recall.
| Model | Class | Precision | Recall | F1 Score | Support |
|---|---|---|---|---|---|
| ResNet-50 | First-Degree Burn | 0.86 | 0.95 | 0.91 | 110 |
| Second-Degree Burn | 0.94 | 0.90 | 0.92 | 107 | |
| Third-Degree Burn | 0.84 | 0.83 | 0.82 | 108 | |
| Macro Avg | 0.88 | 0.88 | 0.88 | 325 | |
| Weighted Avg | 0.88 | 0.88 | 0.88 | 325 | |
| MOBILENET-V3 | First-Degree Burn | 0.90 | 0.90 | 0.90 | 110 |
| Second-Degree Burn | 0.93 | 0.95 | 0.94 | 107 | |
| Third-Degree Burn | 0.87 | 0.84 | 0.85 | 108 | |
| Macro Avg | 0.90 | 0.90 | 0.90 | 325 | |
| Weighted Avg | 0.90 | 0.90 | 0.90 | 325 |
From the table above, we can see the for each burn severity except for the first degree burn, MOBILENET-V3 had a higher recall rate than ResNet-50. The same can be said for other measured metrics including precision and F1 score (ResNet-50 has the edge at recall for first-degree burns and precision for second-degree burns).
One thing to keep in mind is that the test dataset was only 325 images. In the data above the largest percent difference was 5 percent, which for the second-degree burn which would translate to
Therefore, we go to compare model performance against the next metric: confusion matrix. From the confusion matrix we would be able to tell how many false positives or false negatives both models give out on the test data. This is important because if a user has a first degree burn and the model classifies it as a third-degree, the user would be more away and therefore seek medical attention faster.
| ResNet-50 | MOBILENET-V3 |
|---|---|
 |
 |
When comparing the performance on the test dataset, the difference is almost negligible. For example, ResNet-50 correctly classifies
Therefore, the next thing we look at would be false negatives. Each burn degree has different treatment but applying the treatment for a first-degree burn to a second-degree burn can cause the burn to get worse. Below is a calculation for how many burns were classified with a lower degree than it should be for each model.
-
ResNet-50:
$17+0+11=28$ -
MOBILENET-V3:
$11+0+5=16$
This is a difference of
Due to being in the optimal range of false negatives allowed, we chose to deploy MOBILENET-V3 for the application.
To compare the results from other metrics please visit the Google Colabs above.
The biggest limitation with this project is data.
Another major limitation is a lack of spatial awareness. The LLM is only fed the output of the classifier (prediction and confidence) to provide a diagnosis. Therefore, it is unaware where exactly the burn is located. This is especially important as if the burn is on a sensitive area of the body, it may require a different treatment. To get the valid treatment, the user will have to provide the context when chatting. A scenario such as this is where multimodal models would outweight our methodology.
- Integration of a Database: Currently our application depends on using JSON files for each session to store history. On a large scale, this would not be efficient because managing multiple files would easily get complicated.
- Larger & Diverse Dataset: Most likely will require manual labeling due to limited available burn severity datasets. Right now, we were unable to easily choose a model due to the differences in model performance being only 5 images in the worst cases scenario. A more diverse dataset would allow for better generalization on various scenarios. The folder Limit Testing contains some of these extreme cases that were tested on this model.
- Use of Retrieval-Augmented Generation (RAG): The LLM just uses it's pretrained knowledge. For basic first-aid scenarios this is okay. However, for rare cases such as burns on sensitive body parts, the LLM might hallucinate. If Gemini ever hallucinates the user may use the incorrect treatment. To prevent this, we can implement RAG so that we can provide articles from sources such as the Mayo Clinic for Gemini to reference before providing an output.
- Streamlit Latency: Since the app was hosted on Streamlit's free-tier Community Cloud, it could shut down after a period of inactivity, resulting in a cold start and increased initial latency. A potential solution without additional cost is to periodically send requests to keep the app in a hot state, while also caching the model to reduce startup time.
If you are interested in addressing these limitations or adding additional features to this repo, please follow these steps:
-
Fork and Clone the Repo: Fork this repository to your GitHub account, then clone your fork locally:
git clone <your-fork-url> cd BurnSightAI
-
Create a Virtual Environment: Create a virtual environment to isolate the project's dependencies:
python -m venv .venv
Activate the virtual environment:
Windows:
.venv\Scripts\activate
macOS/Linux:
source .venv/bin/activate -
Install Dependencies: Install the required Python packages:
pip install -r requirements.txt
-
Create the
.envFile: Create a.envfile in the root directory of the repository and add the required API keys:GEMINI_API_KEY=your_api_key_here
Do not commit the
.envfile to GitHub. Make sure.envis included in your.gitignore. -
Test the Application Locally: Start the Streamlit application:
streamlit run app.py
Open the local URL provided in the terminal and test the application's existing functionality. Verify that your changes work as expected and do not introduce issues with existing features.
-
Create a New Branch: Create a separate branch for your changes:
git checkout -b <branch-name>
-
Make Your Changes: Implement your feature or fix while following the existing project structure and coding conventions.
-
Test Your Changes Locally: Run the application again and test the specific functionality you modified. If applicable, test different inputs and edge cases to ensure your changes work correctly.
-
Commit and Push Your Changes: Commit your changes and push the branch to your fork:
git add . git commit -m "Describe your changes" git push origin <branch-name>
-
Create a Pull Request: Open a pull request from your branch to the original repository. Include a brief description of the changes you made, why they were necessary, and any relevant testing or considerations.