image similarity

Классификация изображений и поиск сходства

Этот репозиторий содержит сквозное трубопровод для классификации изображений животных в один из 90 видов и выполнение поиска сходства изображений. Проект использует глубокую сверточную нейронную сеть (RESNET-50) в качестве экстрактора функции, в сочетании с пользовательским уровнем встраивания и окончательной головкой классификации. Кроме того, он включает в себя интерфейс Frontend для пользователей для загрузки изображения и получения как прогнозируемого класса, так и визуально сходных изображений из набора данных.

Обзор

• Архитектура модели:

  • Базовая модель: Resnet-50 (предварительно обучен на ImageNet), с удаленным уровнем классификации.
  • Встроенный слой: плотный слой из 256 нейронов, добавленный после экстрактора функции Resnet-50.
  • Голова классификации: последний плотный слой с 90 нейронами, каждый из которых соответствует одному из 90 классов животных.

• Набор данных:
  Набор данных изображения животных - 90 разных животных
  Этот набор данных охватывает широкий спектр видов животных, обеспечивая широкое представление в 90 классах.

• Детали обучения:

  • Уровень обучения: 0,01
  • Эпохи: 10
  • Размер партии: 16
  • Окончательная тренировка потеря: 0,2204

• Показатели производительности:

  • Точность: 0,9706
  • Напомним: 0,9648
  • F1-показатель: 0,9645

Эти метрики указывают на сильную производительность классификации в наборе данных.

Матрица путаницы

Матрица путаницы ( confusion_matrix.png ) включена для визуализации производительности класса к классу.

Поиск сходства изображений

В дополнение к классификации, это решение предоставляет функцию поиска сходства изображений:

1. Извлечение функций:
   Каждое изображение набора данных передается через модель (слой встраивания RESNET-50 + для создания 256-мерного вектора признаков.

2. Хранение базы данных (SQLite):
   Все 256-мерные векторы функций хранятся в базе данных SQLite.

3. Запрос сходства:
   Когда новое изображение загружается, его встраивание вычисляется, и база данных запрашивается, чтобы найти наиболее похожие изображения, основанные на сходстве вектора.

Это позволяет вам не только предсказывать класс загруженного изображения, но и получать визуально похожие изображения из существующего набора данных.

Интерфейс фронта

Простой фронт включен для облегчения взаимодействия:

1. Загрузите изображение: вы предоставляете изображение животного.
2. Получить прогнозируемый класс: система возвращает прогнозируемый класс, если она соответствует одному из 90 известных классов.
3. Просмотреть аналогичные изображения: интерфейс отображает верхние аналогичные изображения из набора данных, позволяя вам исследовать визуально связанные образцы.

Структура репозитория

 .
├── requirements.txt          # Python dependencies
├── README.md                 # This README file
├── .gitignore                # Git ignore file 
├── LICENSE                   # License file for the project                 
├── main.py                   # Entry point for the Flask application
├── model.py                  # Model definition, loading, and inference utilities 
├── similarity.py             # Functions for similarity search and confusion matrix creation
└── train.py                  # Training script for the model
├── templates/
│   ├── index.html            # Frontend template
└── confusion_matrix.png      # Confusion matrix visualization 

Начиная

1. Клонировать репозиторий:

   git clone https://github.com/git-khandelwal/image-similarity
   cd image-similarity

2. Установить зависимости:

   pip install -r requirements.txt

3. Скачать и подготовить набор данных:

   • Скачать с Kaggle: https://www.kaggle.com/datasets/iamsouravbanerjee/animal-image-dataset-90-different-animals
   • Извлекать в базовый каталог.

4. Обучение модели: для переподготовки модели и создания векторных функций DB:

   python train.py

5. Запустите приложение Flask:

   python main.py

6. Используйте приложение:
   • Доступ к приложению по адресу http: // localhost: 5000
   • Загрузите изображение через предоставленную форму.
   • Посмотреть прогнозируемый класс и аналогичные изображения.

Будущие улучшения

• Методы индексации: рассмотрите возможность использования векторных баз данных, таких как FAISS для более быстрых, более масштабируемых поисков сходства.

Ссылки

• К. Хе, Х. Чжан, С. Рен и Дж. Сан, «Глубокое остаточное обучение для распознавания изображений», в «Слушаниях конференции IEEE по компьютерному зрению и распознаванию моделей» (CVPR) , 2016. Paper Link

Лицензия

Этот проект лицензирован по лицензии MIT.