Image Classification – Multi-Class CNN Classifier

Image classification system based on Convolutional Neural Networks (CNN) in TensorFlow/Keras. The project enables training multi-class models and classifying images through a graphical user interface.

Table of Contents

• Features
• Project Structure
• Requirements
• Installation
• Usage
• Model Architecture
• Authors

Features

• Automatic class detection – system scans the classes/ folder and automatically identifies all categories
• Binary and multi-class classification – supports any number of classes (≥2)
• Classification GUI – graphical interface (CustomTkinter) for loading and classifying images
• Data cleaning – automatic removal of corrupted or unsupported image files
• Training monitoring – TensorBoard integration, loss/accuracy plots
• Top-K predictions – display multiple most probable classes with confidence levels

Project Structure

Image-Classification/
├── classes/                    # Training data directory
│   ├── class_1/                # Example: glass_failure, screws_success, etc.
│   ├── class_2/
│   └── class_n/
├── models/
│   └── model.h5                # Trained model (generated)
├── logs/                       # TensorBoard logs
├── testing/                    # Test dataset
├── paper/                      # Research documentation
├── research/                   # Research materials
├── train_model.py              # Model training script
├── image_classifier.py         # Image classification GUI
└── README.md

Requirements

• Python: 3.8+
• Libraries:
  • TensorFlow 2.x
  • OpenCV (opencv-python)
  • CustomTkinter
  • Pillow (PIL)
  • NumPy
  • Matplotlib
  • scikit-learn (optional)

Installation

1. Clone the repository

git clone https://github.com/KarolZygadlo/Image-Classification.git
cd Image-Classification

2. Install dependencies

Option A: Using pip directly

pip install tensorflow opencv-python matplotlib customtkinter pillow numpy

Option B: Using virtual environment (recommended)

python -m venv .venv
source .venv/bin/activate  # On Windows: .venv\Scripts\activate
pip install tensorflow opencv-python matplotlib customtkinter pillow numpy

Note for macOS users with pyenv: If you encounter ModuleNotFoundError: No module named '_tkinter', you need to reinstall Python with Tcl/Tk support:

brew install tcl-tk
env PYTHON_CONFIGURE_OPTS="--with-tcltk-includes='-I/opt/homebrew/opt/tcl-tk/include' --with-tcltk-libs='-L/opt/homebrew/opt/tcl-tk/lib -ltcl9.0 -ltk9.0'" pyenv install <version> --force

3. Prepare data

Create directory structure in the classes/ folder:

classes/
├── class_1/
│   ├── image1.jpg
│   ├── image2.png
│   └── ...
├── class_2/
│   └── ...
└── class_n/
    └── ...

Each subfolder represents a separate class, and its name will be used as a label.

Usage

Training the model

python train_model.py --data-dir classes --epochs 70 --batch-size 32

Available arguments:

• --data-dir – path to folder with classes (default: classes)
• --clean-dir – folder to clean from corrupted images (default: same as --data-dir)
• --img-size – image size (default: 256×256)
• --batch-size – batch size (default: 32)
• --epochs – number of epochs (default: 70)
• --logdir – TensorBoard logs directory (default: logs)
• --model-path – model save path (default: models/model.h5)
• --test-image – path to test image (optional)
• --plot-dir – directory to save plots (optional)
• --show-plots – display plots instead of saving
• --top-k – number of classes to display for test image (default: 3)

Example with all options:

python train_model.py \
  --data-dir classes \
  --epochs 100 \
  --batch-size 16 \
  --plot-dir plots \
  --test-image testing/screws/failure_2.png \
  --top-k 5

Classification via GUI

python image_classifier.py

The interface allows you to:

• Load an image using the "Load Image" button
• See a thumbnail of the image
• Get the predicted class name and confidence level

Note: GUI requires tkinter with Tcl/Tk support. On macOS with pyenv, you need to reinstall Python with --with-tcltk-* flags (see troubleshooting section).

Training monitoring with TensorBoard

tensorboard --logdir logs

Open your browser at http://localhost:6006.

Model Architecture

The model consists of:

1. Convolutional layers:

   • Conv2D(32, 3×3) + ReLU + MaxPooling
   • Conv2D(64, 3×3) + ReLU + MaxPooling
   • Conv2D(32, 3×3) + ReLU + MaxPooling

2. Fully connected layers:

   • Flatten
   • Dense(256, ReLU)
   • Dropout(0.3)
   • Dense(num_classes, Softmax)

3. Loss function: SparseCategoricalCrossentropy

4. Optimizer: Adam

Troubleshooting

No GPU / Low performance

Make sure you have TensorFlow with GPU support (CUDA/cuDNN) installed, or use the CPU version.

Too few batches in dataset

If you see an error about insufficient number of batches, decrease --batch-size or add more images.

License

This project is available under the MIT License.

Authors

• Karol Zygadło – KarolZygadlo

---

Repository: https://github.com/KarolZygadlo/Image-Classification