image similarity

Classificação de imagem e pesquisa de similaridade

Este repositório contém um pipeline de ponta a ponta para classificar imagens de animais em uma das 90 espécies e executar pesquisas de similaridade de imagem. O projeto aproveita uma profunda rede neural convolucional (Resnet-50) como um extrator de recurso, combinado com uma camada de incorporação personalizada e um chefe de classificação final. Além disso, inclui uma interface de front -end para os usuários carregarem uma imagem e receber uma classe prevista e imagens visualmente semelhantes do conjunto de dados.

Visão geral

• Arquitetura de modelo:

  • Modelo básico: Resnet-50 (pré-treinado no ImageNet), com a camada de classificação final removida.
  • Camada de incorporação: uma camada densa de 256 neurônios adicionados após o extrator de recurso Resnet-50.
  • Cabeça de classificação: uma camada densa final com 90 neurônios, cada um correspondente a uma das 90 classes de animais.

• Conjunto de dados:
  Conjunto de dados de imagem animal - 90 animais diferentes
  Esse conjunto de dados cobre uma ampla variedade de espécies animais, garantindo uma ampla representação em 90 classes.

• Detalhes de treinamento:

  • Taxa de aprendizado: 0,01
  • Épocas: 10
  • Tamanho do lote: 16
  • Perda de treinamento final: 0,2204

• Métricas de desempenho:

  • Precisão: 0,9706
  • Lembre -se: 0,9648
  • F1-score: 0,9645

Essas métricas indicam um forte desempenho de classificação no conjunto de dados.

Matriz de confusão

Uma matriz de confusão ( confusion_matrix.png ) está incluída para visualizar o desempenho da classe por classe.

Pesquisa de similaridade da imagem

Além da classificação, esta solução fornece um recurso de pesquisa de similaridade de imagem:

1. Extração de recursos:
   Cada imagem do conjunto de dados é passada através do modelo (camada de incorporação Resnet-50 +) para produzir um vetor de recurso 256-dimensional.

2. Armazenamento de banco de dados (SQLite):
   Todos os vetores de recursos 256-dimensionais são armazenados em um banco de dados SQLite.

3. Consulta de similaridade:
   Quando uma nova imagem é enviada, sua incorporação é calculada e o banco de dados é consultado para encontrar as imagens mais semelhantes com base na similaridade vetorial.

Isso permite que você não apenas preveja a classe de uma imagem carregada, mas também recupere imagens visualmente semelhantes do conjunto de dados existente.

Interface front -end

Um front -end simples está incluído para facilitar a interação:

1. Carregue uma imagem: você fornece uma imagem de um animal.
2. Obtenha classe prevista: o sistema retorna a classe prevista se corresponder a uma das 90 classes conhecidas.
3. Veja imagens semelhantes: a interface exibe as principais imagens semelhantes do conjunto de dados, permitindo explorar amostras relacionadas visualmente.

Estrutura do repositório

 .
├── 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 

Começando

1. Clone o repositório:

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

2. Instale dependências:

   pip install -r requirements.txt

3. Baixe e prepare o conjunto de dados:

   • Download de Kaggle: https://www.kaggle.com/datasets/iamsouravbanerjee/animal-image-dataset-90-different-animals
   • Extraia para o diretório base.

4. Treinamento do modelo: para treinar o modelo e criar db de vetor de recursos:

   python train.py

5. Execute o aplicativo Flask:

   python main.py

6. Use o aplicativo:
   • Acesse o aplicativo em http: // localhost: 5000
   • Carregue uma imagem através do formulário fornecido.
   • Veja a classe prevista e imagens similares.

Melhorias futuras

• Técnicas de indexação: considere o uso de bancos de dados de vetores como o FAISS para pesquisas de similaridade mais rápidas e escaláveis.

Referências

• K. He, X. Zhang, S. Ren e J. Sun, "Aprendizagem residual profunda para reconhecimento de imagem", em Anais da Conferência IEEE sobre Visão Computador e Reconhecimento de Padrões (CVPR) , 2016. Link de papel

Licença

Este projeto está licenciado sob a licença do MIT.