Brain Tumor Classifier

Project Explanation

This project used a Neural Network to classify different kinds of brain tumors. Given an image of a brain MRI scan, it would either predict that the brain contains no tumors or classify the tumors as either glioma, meningioma or pituitary. I later created a user interface for the model.

Git Hub Repo link:
https://github.com/Thomasche69/CNNBraintumor

Training dataset

The training data was organized into this directory structure. It was divided into two main folders, one for training the model and one for testing. Each folder contained all the images for the tumor types glioma, meningioma and pituitary and images without tumors all seperated into different folders of their own. The dataset was overall very balanced with each separate type of tumor and no tumors having roughly 1300 images each, making the dataset very balanced.

No Tumor Brain

glioma

Model architecture

The Neural Network architecture used in this project is a Convolutional Neural Network (CNN). Further information of CNNs can be found in the AI Concepts section of this website. Other architectures such as Transformers and pretrained Vision Models were attempted on this project, however this architecture yielded the best result.

CNN class initialization code

forward pass code:

the comments in the code show the tensor dimensions of x before each layer. For example the first one is [batch_size, (Red, Green, Blue), height of image, width of image]

Preprocessing before Training

Before the data was fed into the model, preprocessing was done in order to increase the size of the dataset as well as having the model learn from noisier data. The transformations for the training data and test data was different. The training part applied random flips, vertical flips, and rotations of 10 degrees to generate more images in the training dataset. This allows the model to learn from more unusual inputs and noise, which would increase its performance on inference. The testing data only applied resizing the images to standardize model inputs as well as applying a normalization to reduce the spread of the data. This normalization ensures that the model doesn't overfit when training.

This function simply loads the dataset and applies the transformations to each image. The data loaders shuffle the data and organize them into batches of size 32. Meaning that one batch contains 32 images.

Training loop

The criterion used for evaluating the model's performance was CrossEntropyLoss while Adam was the main optimizer used during the training.

One epoch meant going through the entire dataset once. Meaning if there was 10 epochs, the model would go through the training data 10 times.

During each epoch the model would loop through the entirety of the training data. First it would receive the input images from the data loaders. Then the output of that model would be evaluated using CrossEntropyLoss. The optimizer will reset all the gradients to 0, then using the loss function calculate the gradients of the model via back propagation, then take one step to update the model's weights.

The loss values are summed up throughout the training loop and then divided by the size of the training dataset

It is then evaluated by running through the entire testing dataset. The loss is calculated the same way as in the training dataset. This process ends once the model has gone through the designated number of epochs

The model was set to have 40 epochs and the learning rate given to the optimizer was 0.0002. The model was then saved after training

Results

The model ended up achieving a 98.8% accuracy in predicting the presence of tumor types. on the testing dataset Meanwhile the training accuracy achieved a slightly higher result of 98.9% accuracy, indicating that the model didn't experience significant underfitting or overfitting. Based on the graph in figure 1, the overall model's performance on the validation set was unstable in earlier epochs but become much more stable in later epochs.

Figure 1: Training and validation accuracy of the model throughout 40 epochs

User interface demo

After the model was saved after training, I then hosted this model on a streamlit website with a simple interface. This website also has the feature of an LLM explanation of the type of cancer that was detected.

Screen Recording 2025-10-11 at 01.48.36.mov

APIs

The APIs were made using flask and hosted using waitress.

This predict endpoint is responsible for hosting the CNN model and providing a prediction based on the uploaded image. It first extracts the the image from the requested files. This image is then saved to a temporary directory. These images are then preprocessed and fed to the model which provides the predictions and probabilities. These outputs are then converted from tensors to lists and integers, for it to be returned via json. This endpoint then returns this json containing the predicted tumor class or no tumor and the probability distributions.

The processing function using the test transforms from earlier. However this image is only a 3D tensor [channels, height, width] but the model expects a 4D tensor [batch_size, channels, height, width]. The image simply unsqueezes an extra dimension [1, channels, height, width] to solve this issue.

This image after preprocessing is then fed to the model. The model then returns an output which is a probability distribution of what it is most confident in each class. For example it might have high confidence in glioma and low confidence that the image is not a tumor. The torch.max function then selects the class with the highest confidence, making that type of cancer or not its prediction.

This endpoint is much simpler, it simply recieves the predicted tumor type of the previous model and passes it into a prompt of the LLM. The prompt asks the model to explain the tumor type in simple terms. After passing the prompt to the model, the model response is then returned in a json file.

Streamlit app code

This is the main code that runs the user interface and communicates with the APIs. The uploaded_file part is responsile for handling the process of uploading images. This uploaded image is then passed to the /predict endpoint to then recieve the CNN's prediction. After the prediction is recieved, the website then writes the predicted tumor type. If there is no predicted tumor, then the website will simply write that no tumor was detected. If there was a tumor type predicted, the name of the tumor is then passed to the LLM endpoint.

The LLM is a reasoning model, meaning that its response is divided into two main parts a "thinking" part and the main text part. The "thinking" part helps the LLM formulate a better response however we only need the main part for our website. So using regex, the text surrounded by the think blocks are seperated from the main texts. The main texts is then displayed while the thinking part is removed.

Page updated

Report abuse