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許可獲得許可。