Automatic Traffic Sign Classification

Проект для классификации дорожных знаков с использованием глубокого обучения. Обучает модель ResNet18 на датасете GTSRB (German Traffic Sign Recognition Benchmark) для распознавания 43 классов дорожных знаков.

📋 Содержание

• Особенности
• Требования
• Установка
• Быстрый старт
• Структура проекта
• Конфигурация
• Использование
• Разработка
• Лицензия

✨ Особенности

• 🚀 PyTorch Lightning — современный фреймворк для обучения моделей
• 🔧 Hydra — гибкая система управления конфигурациями
• 📊 MLflow — отслеживание экспериментов и метрик
• 📦 DVC — версионирование данных
• 🎯 ResNet18 — предобученная архитектура для классификации изображений
• 📥 Kaggle API — автоматическая загрузка датасета GTSRB

🔧 Требования

• Python >= 3.13
• CUDA-capable GPU (опционально, для ускорения обучения)
• Kaggle API credentials (для загрузки датасета)

Setup (📦 Установка)

1. Клонирование репозитория

git clone https://github.com/ek2kk/automatic-traffic-sign-classification
cd automatic-traffic-sign-classification

2. Установка зависимостей

Проект использует uv для управления зависимостями:

uv sync

Или с помощью pip:

pip install -e .

3. Настройка Kaggle API

Для загрузки датасета необходимо настроить Kaggle API:

1. Создайте API токен:

   • Зайдите на Kaggle.com
   • Перейдите в Settings → API → "Create New Token"
   • Скачайте файл kaggle.json

2. Разместите токен:

   mkdir -p ~/.kaggle
   cp kaggle.json ~/.kaggle/kaggle.json
   chmod 600 ~/.kaggle/kaggle.json

   Альтернатива: Используйте переменные окружения:

   export KAGGLE_USERNAME="your_username"
   export KAGGLE_KEY="your_api_key"

Train (🚀 Быстрый старт)

1. Загрузка данных

uv run python -m automatic_traffic_sign_classification.data.download

Данные будут скачаны в data/gtsrb/ (~200MB). Процесс может занять несколько минут.

2. Обучение модели

uv run python -m automatic_traffic_sign_classification.commands

Модель начнет обучение с параметрами из configs/train.yaml.

3. Использование DVC (опционально)

Если используете DVC для версионирования данных:

dvc repro download_data

📁 Структура проекта

automatic-traffic-sign-classification/
├── automatic_traffic_sign_classification/    # Основной пакет
│   ├── __init__.py
│   ├── commands.py                          # Точка входа (Hydra)
│   ├── train.py                             # Логика обучения
│   ├── data/
│   │   ├── __init__.py
│   │   ├── dataset.py                      # PyTorch Dataset для GTSRB
│   │   └── download.py                     # Загрузка данных с Kaggle
│   ├── models/
│   │   ├── __init__.py
│   │   └── classifier.py                   # ResNet18 классификатор
│   ├── lightning/
│   │   ├── __init__.py
│   │   └── module.py                       # Lightning модуль
│   └── utils/
│       ├── __init__.py
│       └── seed.py                         # Фиксация seed для воспроизводимости
├── configs/                                 # Конфигурационные файлы (Hydra)
│   ├── train.yaml                          # Главный конфиг
│   ├── data/
│   │   ├── base.yaml                      # Базовый конфиг данных
│   │   └── kaggle.yaml                    # Конфиг для Kaggle датасета
│   ├── model/
│   │   └── resnet18.yaml                  # Параметры модели
│   ├── trainer/
│   │   └── default.yaml                   # Параметры тренера
│   └── logging/
│       └── mlflow.yaml                     # Настройки MLflow
├── data/                                    # Данные (не в git)
│   └── gtsrb/                              # GTSRB датасет
├── pyproject.toml                          # Зависимости проекта
├── dvc.yaml                                # DVC pipeline
├── docker-compose.yml                      # Docker конфигурация
└── README.md                               # Этот файл

⚙️ Конфигурация

Проект использует Hydra для управления конфигурациями. Все параметры находятся в configs/.

Основные конфигурационные файлы

configs/train.yaml — главный конфиг

defaults:
  - model/resnet18
  - data/kaggle
  - trainer/default
  - logging/mlflow

train:
  batch_size: 128
  lr: 3e-4
  use_gpu: true

configs/data/kaggle.yaml — параметры данных

data_dir: ./data/gtsrb/Train
batch_size: 128
val_split: 0.2
num_workers: 4
image_size: 32

configs/model/resnet18.yaml — параметры модели

num_classes: 43
pretrained: true

configs/trainer/default.yaml — параметры тренера

max_epochs: 10
precision: 32
gradient_clip_val: 0.0

Переопределение параметров

Вы можете переопределить любые параметры через командную строку:

uv run python -m automatic_traffic_sign_classification.commands \
  train.batch_size=64 \
  train.lr=1e-4 \
  trainer.max_epochs=20 \
  train.use_gpu=false

Или создать новый конфиг в configs/ и указать его:

uv run python -m automatic_traffic_sign_classification.commands \
  --config-name=my_config

📊 Использование

Загрузка данных

# Стандартная загрузка
uv run python -m automatic_traffic_sign_classification.data.download

# Или через DVC
dvc repro download_data

Структура данных после загрузки:

data/gtsrb/
└── Train/
    ├── 0/          # Класс 0
    │   ├── 00000.ppm
    │   ├── 00001.ppm
    │   └── ...
    ├── 1/          # Класс 1
    └── ...

Датасет содержит 43 класса дорожных знаков, каждый класс находится в отдельной директории.

Обучение модели

# Базовое обучение
uv run python -m automatic_traffic_sign_classification.commands

# С переопределением параметров
uv run python -m automatic_traffic_sign_classification.commands \
  train.batch_size=256 \
  trainer.max_epochs=50

# Обучение на CPU
uv run python -m automatic_traffic_sign_classification.commands \
  train.use_gpu=false

Отслеживание экспериментов (MLflow)

Для использования MLflow убедитесь, что сервер запущен:

mlflow ui --port 8080

Затем запустите обучение — метрики будут автоматически логироваться в MLflow.

🔬 Разработка

Структура кода

• data/dataset.py — реализует GTSRBDataset, PyTorch Dataset для загрузки изображений
• data/download.py — загрузка датасета GTSRB с Kaggle
• models/classifier.py — архитектура ResNet18 с заменой последнего слоя
• lightning/module.py — Lightning модуль с логикой обучения и оптимизатором
• train.py — основная функция обучения, создает DataLoader и Trainer
• commands.py — точка входа с Hydra декоратором

Добавление новых моделей

1. Создайте новый файл в automatic_traffic_sign_classification/models/
2. Реализуйте класс модели, наследуясь от nn.Module
3. Создайте конфиг в configs/model/
4. Обновите train.py для использования новой модели

Добавление валидации

Текущая версия обучает только на тренировочных данных. Для добавления валидации:

1. Добавьте validation_step в lightning/module.py
2. Разделите датасет на train/val в train.py
3. Передайте оба DataLoader в trainer.fit()

Тестирование

# Проверка импортов
python -c "from automatic_traffic_sign_classification.models.classifier import TrafficSignClassifier; print('OK')"

# Проверка загрузки данных (после скачивания)
python -c "from automatic_traffic_sign_classification.data.dataset import GTSRBDataset; ds = GTSRBDataset('data/gtsrb/Train'); print(f'Dataset size: {len(ds)}')"

📈 Производительность

• Архитектура: ResNet18 (предобученная на ImageNet)
• Размер изображений: 32x32 пикселей
• Классы: 43 класса дорожных знаков
• Оптимизатор: Adam
• Loss: Cross Entropy

Ожидаемая точность на валидационной выборке: ~95-98% после 10-20 эпох обучения.

Автор: Egor Kuznetsov Версия: 0.1.0