image similarity

图像分类和相似性搜索

该存储库包含一种将动物图像分类为90种的端到端管道，并执行图像相似性搜索。该项目利用深卷积神经网络（Resnet-50）作为功能提取器，并结合自定义嵌入层和最终分类头。此外，它包括一个前端接口，供用户上传图像并从数据集中接收预测的类和视觉上相似的图像。

概述

• 模型体系结构：

  • 基本模型： Resnet-50（在ImageNet上进行了预训练），最终分类层已删除。
  • 嵌入层： Resnet-50提取器后添加的256个神经元的致密层。
  • 分类头：最终致密层，带有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. 运行烧瓶应用程序：

   python main.py

6. 使用该应用程序：
   • http：// localhost访问该应用程序：5000
   • 通过提供的表单上传图像。
   • 查看预测的类和类似的图像。

未来的改进

• 索引技术：考虑使用Faiss等矢量数据库，以更快，更可扩展的相似性搜索。

参考

• K. He，X。Zhang，S。Ren和J. Sun，“图像识别的深度残留学习”，在IEEE计算机视觉和模式识别会议论文集（CVPR） ，2016年。纸质链接

执照

该项目已根据MIT许可获得许可。