CellVision AI

A computer vision project for classifying white blood cells, detecting unusual images, and explaining model predictions.

I started this project as a simple image classification task, but later added transfer learning, anomaly detection with a VAE, Grad-CAM, and a small segmentation experiment with U-Net.

The project currently works with four blood cell classes:

• Eosinophil
• Lymphocyte
• Monocyte
• Neutrophil

This is an educational project, not a medical diagnostic tool.

---

What I did

The main steps of the project were:

1. explored the dataset and checked class balance;
2. trained a simple CNN as a baseline;
3. improved the CNN with augmentation and regularization;
4. tested ResNet-18 and EfficientNet-B0 with transfer learning;
5. trained a VAE for anomaly detection;
6. combined the classifier and VAE into one hybrid pipeline;
7. added Grad-CAM to see what image regions affected the prediction;
8. made a small U-Net segmentation experiment using automatically generated masks;
9. added a Streamlit interface and deployment files.

---

Dataset

The dataset contains microscope images of four white blood cell types.

Class	Training images
Eosinophil	2,497
Lymphocyte	2,483
Monocyte	2,478
Neutrophil	2,499

Dataset split used in the notebook:

Split	Images
Full training source	9,957
Training subset	7,965
Validation subset	1,992
Test set	2,487

The dataset is almost perfectly balanced, so I did not use class weights or oversampling.

The original image size is around 240 x 320, but I resized images for training:

• 128 x 128 for classification;
• 64 x 64 for the VAE.

The dataset itself is not included in this repository.

---

Preprocessing

For the classification models I used:

• resize;
• tensor conversion;
• normalization;
• random horizontal flip;
• random vertical flip;
• random rotation;
• other basic augmentation.

The main reason for augmentation was to reduce overfitting. The first CNN was learning the training set quite well, but its test result was much worse.

---

Models

Baseline CNN

The first model was a simple CNN with convolution, ReLU, max pooling, and fully connected layers.

It gave me a starting point, but it overfitted and confused visually similar classes, especially Eosinophil and Neutrophil.

Improved CNN

After that I added:

• more convolutional layers;
• batch normalization;
• dropout;
• stronger augmentation.

This improved test accuracy from about 68% to 87%.

ResNet-18

I tested two options:

• frozen pretrained backbone;
• full fine-tuning.

The frozen version did not work very well, while fine-tuning gave a much better result.

EfficientNet-B0

EfficientNet-B0 was also tested with a frozen backbone and with full fine-tuning.

The fine-tuned version gave the best validation result, so I used it as the final classifier.

VAE for anomaly detection

I trained a convolutional Variational Autoencoder on the known blood cell images.

The idea is simple: if the VAE cannot reconstruct an input image well, the image may be unusual or outside the training distribution.

The anomaly threshold was calculated as:

threshold = mean reconstruction error + 3 * standard deviation

Values from the notebook:

Statistic	Value
Mean reconstruction error	0.0137
Standard deviation	0.0023
Threshold	0.0207

---

Hybrid model

The final logic is:

Input image
    |
    v
VAE reconstruction
    |
    |-- error > threshold --> anomaly
    |
    |-- error <= threshold
                |
                v
        EfficientNet-B0
                |
                v
      class + confidence

So the model does not always force an image into one of the four classes. First, it checks whether the image looks similar to the data used during training.

---

Grad-CAM

I used Grad-CAM to visualize the image regions that influenced the final prediction.

The notebook shows:

• original image;
• heatmap;
• heatmap overlay;
• predicted class;
• confidence score.

This part was useful for checking whether the classifier was looking at the cell area and not only at the background.

---

Segmentation experiment

I also made a small U-Net segmentation experiment.

There were no manually labeled masks, so I generated pseudo-masks with:

• grayscale conversion;
• Otsu thresholding;
• morphological opening;
• morphological closing.

The U-Net was trained with BCE and Dice loss.

The result was very high, but this should be interpreted carefully because the model was learning to reproduce automatically generated masks, not expert annotations.

---

Results

Model	Split	Accuracy	Macro F1
Baseline CNN	Test	0.680	0.690
Improved CNN	Test	0.870	0.870
ResNet-18, frozen	Validation	0.550	0.550
ResNet-18, fine-tuned	Validation	0.978	0.980
EfficientNet-B0, frozen	Validation	0.754	0.750
EfficientNet-B0, fine-tuned	Validation	0.998	1.000

U-Net experiment:

Metric	Value
Training samples	800
Validation samples	200
Best validation Dice	0.9990

The EfficientNet result is very high, so I would not treat it as final proof of real-world performance. It still needs a strict external test set and a duplicate check.

---

Project structure

cellvision-ai/
├── Dockerfile
├── best_model_efficientnet.pth
├── hybrid_model.py
├── hybrid_model_config.json
├── main.py
├── project.ipynb
├── railway.toml
├── requirements.txt
├── streamlit_app.py
├── vae_anomaly.pth
└── README.md

Main files:

• project.ipynb — EDA, training, evaluation, Grad-CAM, and segmentation experiments
• best_model_efficientnet.pth — trained classifier
• vae_anomaly.pth — trained VAE
• hybrid_model.py — inference logic
• hybrid_model_config.json — class names and anomaly threshold
• streamlit_app.py — Streamlit interface
• Dockerfile — Docker configuration
• railway.toml — Railway deployment configuration

---

How to run

Clone the repository:

git clone https://github.com/weakyweaky7/cellvision-ai.git
cd cellvision-ai

Install dependencies:

pip install -r requirements.txt

Run the Streamlit app:

streamlit run streamlit_app.py

Or open the notebook:

jupyter notebook project.ipynb

---

Docker

Build the image:

docker build -t cellvision-ai .

Run it:

docker run --rm -p 8501:8501 cellvision-ai

Then open:

http://localhost:8501

---

Tech stack

• Python
• PyTorch
• Torchvision
• OpenCV
• NumPy
• pandas
• Matplotlib
• Seaborn
• scikit-learn
• Streamlit
• Docker
• Railway

---

Limitations

• EfficientNet was evaluated on the validation split, not on an external dataset.
• The VAE threshold is experimental and was not tested on a labeled anomaly dataset.
• U-Net was trained on pseudo-masks, not on manually annotated masks.
• Model confidence is not calibrated probability.
• Results may change on images from another microscope or laboratory.
• The dataset should be checked for duplicate and near-duplicate images.

---

Next steps

• test the classifier on external data;
• check train and validation images for duplicates;
• calibrate probabilities;
• evaluate anomaly detection on real out-of-distribution examples;
• use manually annotated masks for segmentation;
• split training and inference code into separate modules;
• add tests and experiment tracking.

---

Disclaimer

This repository was made for learning and portfolio purposes. It should not be used for diagnosis or any real medical decision.