An end-to-end deep-learning project that classifies gallery images into five categories:

blurred · Notes · Posters · real · Screenshots

This repository focuses only on image-model training: dataset preparation, model training, evaluation, TensorFlow Lite export, and single-image prediction.

About

Phone galleries contain normal photos, screenshots, notes, posters, and accidental blurry images. This project trains a supervised image classifier that identifies those categories automatically.

The model uses transfer learning with EfficientNetV2B0. It learns from labelled images, is tested on unseen validation images, and exports a .tflite model after training.

Dataset structure

Place images in class folders as shown below:

data/
├── blurred/
├── Notes/
├── Posters/
├── real/
└── Screenshots/

Class	Images	Description
blurred	700	Blurry or low-quality photos
Notes	960	Notes, documents, and written reminders
Posters	896	Posters and graphic images
real	998	Normal photographs
Screenshots	741	Screen captures

Total images: 4,295

Do not upload personal images to a public GitHub repository unless you have permission to share them.

Workflow

Labelled images
      ↓
prepare_dataset.py
      ↓
Clean RGB JPEG images
      ↓
train.py
      ↓
EfficientNetV2B0 model
      ↓
Metrics + Keras model + TFLite model
      ↓
predict.py
      ↓
Prediction for a new image

Features

• Supports JPG, JPEG, PNG, WEBP, GIF, and BMP inputs
• Corrects image orientation and converts images to RGB JPEG
• Uses an 80/20 training-validation split
• Uses a fixed seed for reproducible results
• Uses EfficientNetV2B0 transfer learning
• Applies augmentation: flip, rotation, zoom, and contrast changes
• Uses early stopping and saves the best model
• Exports a TensorFlow Lite model
• Includes Google Colab GPU training support

Files

File	Purpose
prepare_dataset.py	Converts source images into consistent RGB JPEG training data.
train.py	Trains, validates, and exports the classifier.
predict.py	Predicts the class of one image using the trained TFLite model.
colab_train.ipynb	Google Colab notebook for GPU training.
requirements.txt	Required Python packages.

Installation

git clone <your-repository-url>
cd IMAGE_classifier
python -m venv .venv

Activate the environment.

Windows PowerShell

.\.venv\Scripts\Activate.ps1

macOS / Linux

source .venv/bin/activate

Install requirements:

pip install -r requirements.txt

Training locally

1. Prepare images

python prepare_dataset.py --source data --output processed_data

This creates a cleaned processed_data/ folder. Your original images in data/ are not changed.

2. Train the classifier

python train.py --data processed_data --output artifacts --epochs 15 --batch-size 32

Training uses 80% of images to learn and 20% for validation.

Train on Google Colab

Use Google Colab if your computer cannot handle model training.

1. Upload the complete project folder to Google Drive.
2. Open colab_train.ipynb in Google Colab.
3. Select Runtime → Change runtime type → T4 GPU.
4. Run the cells from top to bottom.
5. The finished files are saved in artifacts/.

Output files

After training, the artifacts/ folder contains:

artifacts/
├── best.keras
├── model.keras
├── gallery_cleanup_classifier.tflite
├── labels.txt
├── metrics.json
└── training_config.json

File	Meaning
model.keras	Full TensorFlow/Keras model.
best.keras	Best model based on validation accuracy.
gallery_cleanup_classifier.tflite	Lightweight model for inference.
labels.txt	Class names in model-output order.
metrics.json	Validation accuracy and training history.

Test one image

After training, run:

python predict.py "data/Screenshots/your-image.png"

Example output:

Image: data/Screenshots/your-image.png
Screenshots  95.4%
real          2.3%
Notes         1.1%
blurred       0.8%
Posters       0.4%

The class with the highest confidence is the prediction.

Evaluation

Open artifacts/metrics.json after training. The first metric to check is:

validation.accuracy

Also test images that were not part of the original dataset. Accuracy alone is not enough; check whether the model confuses categories such as normal photos and blurry images.

Technology used

• Python
• TensorFlow / Keras
• EfficientNetV2B0
• TensorFlow Lite
• Pillow
• NumPy
• Google Colab

Future improvements

• Add a confusion matrix and per-class precision/recall.
• Add more images for weaker classes.
• Fine-tune the EfficientNet backbone after initial training.
• Test with images from different phones and lighting conditions.