Skip to content

baldellimtt/tensorflow-image-classifier

Repository files navigation

Image Classification with TensorFlow

Binary image classifier with transfer learning (MobileNet), YAML configuration, and GPU support on WSL2/Ubuntu via uv.

Project Structure

├── train.py              # Training entrypoint
├── evaluate.py           # Test-set evaluation and threshold search
├── multi_train.py        # Run multiple config files sequentially
├── config.yaml           # Default training configuration
├── model/                # Model architecture
├── utils/                # Config, data split, training pipeline, GPU setup
├── scripts/              # WSL setup and GPU check helpers
└── pyproject.toml        # Dependencies (uv)

Setup (WSL / Ubuntu)

cd /mnt/c/Users/Matteo\ Baldelli/Desktop/tensorflow-image-classifier
bash scripts/setup-wsl.sh

Or manually:

export UV_LINK_MODE=copy
export UV_PROJECT_ENVIRONMENT="$HOME/.venvs/tensorflow-image-classifier"
uv sync
uv run python scripts/check_gpu.py

Training

uv run python train.py
uv run python train.py --config path/to/config.yaml

Dataset layout:

data/
  class_0/
  class_1/

The split is written to tmp_split/{train,val,test} on each run. Logs and metrics go to logs/<run_name>/.

Evaluation

uv run python evaluate.py -m modelli/<model_path> [--threshold 0.5] [--mode standard|folder_split]

Configuration

config.yaml is validated at startup with Pydantic. Key fields:

Field Description
auto_split Source/output dirs and val/test ratios
input_shape [H, W, 3] — images are resized to 224×224 inside MobileNet
batch_size, epochs, lr Training hyperparameters
decay_rate, decay_steps Exponential learning-rate decay
seed Reproducibility seed
checkpoint_filepath Where best models are saved

Development

uv sync --group dev
uv run ruff check .

About

Config-driven TensorFlow/Keras image classification pipeline — transfer learning, YAML-based training, custom F1 metrics, threshold tuning, and model evaluation.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages