How arshad1428/image-classification-cnn works
This project trains and serves six-class scene image classifiers. The main workflow loads and splits Intel images, applies training/evaluation transforms, trains a selected CNN, validates and checkpoints it, then evaluates on the held-out test set while recording metrics in MLflow. A separate Streamlit workflow loads a checkpoint and predicts a class and confidence for an uploaded image. The README documents both model variants and experiment tracking; unsampled implementation details are represented only where the documented capability and tree support them.
Main entry points: main.py Training entry · app.py Streamlit app · transfer_learning.py ResNet18 model · predict.py Image prediction
Components 16
Image data
- Intel images
data/raw - Dataset splits
src/image_classifier/data_loader.py - Image transforms
src/image_classifier/preprocessing.py
Training workflow
- Training entry
main.py - Train and validate
src/image_classifier/training/train.py - Project config
src/image_classifier/config.py
Models and artifacts
- Custom CNN
src/image_classifier/models/cnn.py - ResNet18 model
src/image_classifier/models/transfer_learning.py - Checkpoint handling
src/image_classifier/training/checkpoint.py - Saved checkpoints
models
Evaluation and serving
- Test metrics
src/image_classifier/evaluation/evaluate.py - MLflow tracking
- Streamlit app
app.py - Image prediction
src/image_classifier/inference/predict.py
Other components
- User
- Torchvision weights
Connections 29
- User → Streamlit app: uploads image
- Intel images → Dataset splits: reads images
- Dataset splits → Image transforms: gets transforms
- Dataset splits → Training entry: returns loaders
- Project config → Dataset splits: supplies paths
- Project config → Training entry: supplies settings
- Training entry → Dataset splits: creates loaders
- Training entry → Custom CNN: constructs when selected
- Training entry → ResNet18 model: constructs when selected
- ResNet18 model → Torchvision weights: loads weights
- Training entry → Train and validate: runs epochs
- Train and validate → Custom CNN: calls model
- Train and validate → ResNet18 model: calls model
- Training entry → Checkpoint handling: saves best model
- Checkpoint handling → Saved checkpoints: writes checkpoint
- Training entry → Test metrics: evaluates test set
- Test metrics → Custom CNN: calls model
- Test metrics → ResNet18 model: calls model
- Training entry → MLflow tracking: logs params and metrics
- Training entry → MLflow tracking: logs checkpoint artifact
- Streamlit app → Custom CNN: constructs when selected
- Streamlit app → ResNet18 model: constructs when selected
- Streamlit app → Saved checkpoints: loads checkpoint
- Streamlit app → Image prediction: requests prediction
- Image prediction → Image transforms: gets eval transform
- Image prediction → Custom CNN: calls model
- Image prediction → ResNet18 model: calls model
- Image prediction → Streamlit app: returns class and confidence
- Streamlit app → User: shows result