A deep learning project to classify chest X-ray images into three categories: Normal, Viral Pneumonia, and COVID-19 using ResNet18 architecture with PyTorch.
This project demonstrates a complete machine learning pipeline for medical image classification. The model uses transfer learning with a pre-trained ResNet18 architecture to classify chest X-ray images into three categories:
- Normal: Healthy chest X-rays
- Viral: Chest X-rays showing viral pneumonia
- COVID-19: Chest X-rays showing COVID-19 symptoms
The project leverages PyTorch's computational efficiency and includes data preprocessing, model training, validation, and visualization components.
The notebook is divided into the following key parts:
- Imports necessary libraries: PyTorch, torchvision, NumPy, PIL, and Matplotlib
- Sets random seeds for reproducibility
- Displays PyTorch version information
- Organizes raw X-ray images into structured directories
- Renames folders to standardized class names:
normal,viral,covid - Creates a separate test directory with 30 images per class for evaluation
- Uses random sampling to split training and test data
- Implements
ChestXRayDatasetclass extendingtorch.utils.data.Dataset - Loads images from class directories
- Handles image retrieval and random sampling
- Converts images to RGB format for compatibility
- Provides dataset length and item indexing
- Training Transform: Resizes images to 224×224, applies random horizontal flips, converts to tensors, and normalizes
- Test Transform: Resizes to 224×224, converts to tensors, and normalizes
- Uses ImageNet normalization standards (mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
- Creates training dataset pointing to the main image directories
- Creates test dataset pointing to the test subdirectories
- Loads both datasets with the defined transformations
- Creates PyTorch DataLoaders for batching
- Batch size set to 6
- Enables shuffling for training data
- Provides dataset statistics (number of batches)
- Implements
show_images()function to visualize X-ray samples - Displays predicted vs actual labels
- Color-codes correct predictions (green) vs incorrect ones (red)
- Denormalizes images for proper visualization
- Uses pre-trained ResNet18 from torchvision
- Modifies final fully connected layer from 1000 to 3 outputs (one per class)
- Defines Cross-Entropy Loss function for multi-class classification
- Sets up Adam optimizer with learning rate 3e-5
- Implements comprehensive training function with:
- Epoch-based training
- Validation at regular intervals (every 20 steps)
- Loss calculation and backpropagation
- Accuracy metrics computation
- Early stopping when accuracy ≥ 95%
- Training vs evaluation mode switching
- Runs training loop for specified number of epochs
- Visualizes predictions on test data before and after training
- Displays performance improvements over time
- Python 3.7 or higher
- pip or conda package manager
- GPU recommended (optional, for faster training)
git clone https://github.com/EarnTHYPart/COVID-19-Chest-XRay.git
cd COVID-19-Chest-XRayUsing venv:
python -m venv venv
# On Windows
venv\Scripts\activate
# On macOS/Linux
source venv/bin/activateUsing conda:
conda create -n covid-xray python=3.8
conda activate covid-xraypip install torch torchvision matplotlib pillow numpyOr install from requirements.txt:
pip install -r requirements.txtFor faster training on NVIDIA GPUs:
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118The project uses the COVID-19 Radiography Database available on Kaggle.
Download Link: COVID-19 Radiography Database
After downloading, extract the dataset into the project directory:
COVID-19-Chest-XRay/
├── COVID-19 Radiography Database/
│ ├── normal/ (NORMAL images)
│ ├── viral/ (Viral Pneumonia images)
│ ├── covid/ (COVID-19 images)
│ └── test/ (Test split)
│ ├── normal/
│ ├── viral/
│ └── covid/
├── Complete Notebook.ipynb
└── README.md
The notebook automatically reorganizes the downloaded data into the correct structure when first run.
-
Launch Jupyter:
jupyter notebook
-
Open the Notebook:
- Navigate to
Complete Notebook.ipynb - Click to open
- Navigate to
-
Run Cells:
- Click on each cell and press
Shift+Enterto execute - Or use
Kernel → Run Allto execute all cells - Cells run sequentially and depend on previous cells
- Click on each cell and press
jupyter lab- Install the Jupyter extension in VS Code
- Open
Complete Notebook.ipynb - Click the Run button on each cell or use the Jupyter interface
To modify training parameters, edit these variables in the notebook:
# Adjust batch size (Cell 7)
batch_size = 6 # Increase for faster training, decrease for lower memory usage
# Adjust learning rate (Cell 12)
optimizer = torch.optim.Adam(resnet18.parameters(), lr=3e-5)
# Adjust training epochs (Cell 16)
train(epochs=1) # Increase for more training iterationsResNet18 is a 18-layer residual neural network:
- Input: 224×224 RGB images
- Feature Extraction: Multiple residual blocks with skip connections
- Backbone: Pre-trained on ImageNet
- Output Layer: 3 neurons (one per class)
- Final Activation: Softmax (applied by CrossEntropyLoss)
ResNet18 (pretrained)
└── Final FC Layer: 512 → 3 outputs
- Loss Function: CrossEntropyLoss
- Optimizer: Adam (lr=3e-5)
- Batch Size: 6
- Validation Accuracy Target: ≥ 95%
The model achieves:
- High accuracy in distinguishing between normal, viral pneumonia, and COVID-19 X-rays
- Early stopping at ≥95% accuracy for efficient training
- Visualization of predictions with color-coded correctness (green=correct, red=incorrect)
- Validation loss and accuracy are evaluated every 20 training steps
- Sample predictions are visualized during training to monitor performance
- Training typically converges within 1 epoch due to transfer learning
- First Run: The first execution will reorganize the dataset. This is normal and only happens once.
- GPU Memory: If you run out of memory, reduce
batch_sizefrom 6 to 4 or 2 - Training Time: With GPU, training typically takes 10-30 minutes per epoch
- Predictions: Use
show_preds()function anytime to visualize model predictions on test data
| Issue | Solution |
|---|---|
| Dataset not found | Ensure the COVID-19 Radiography Database folder is extracted in the project directory |
| Out of memory | Reduce batch size in Cell 7 |
| Slow training | Enable GPU support or reduce dataset size |
| Import errors | Run pip install -r requirements.txt |
This project is provided as-is for educational and research purposes.
If you use this project, please cite the original dataset:
Rahman, T., Chowdhury, A., Khandakar, A. (2021). COVID-19 Radiography Database. Mendeley Data.
For issues, questions, or contributions, please open an issue on the GitHub repository.
Happy Learning! 🚀
New convenience cells have been added at the end of the notebook:
- Evaluation Metrics: Computes overall test accuracy, a confusion matrix, and per-class precision/recall. It will also render a confusion matrix plot for quick visual inspection.
- Model Save/Load + TorchScript: Saves the trained
resnet18weights tomodels/resnet18_covid_xray.pth, demonstrates loading them into a fresh model, and optionally exports a TorchScript module tomodels/resnet18_covid_xray_script.ptfor deployment.
- Run all training cells as usual.
- Execute the new "Evaluation" cell to see metrics and the confusion matrix.
- Execute the "Save and Load Model" cell to persist weights and produce a TorchScript artifact.
Artifacts will be written under the models/ folder, which is created automatically if it doesn't exist.