找回密码
 立即注册
首页 业界区 业界 撸一个功能强大的基于语义的图像检索系统 ...

撸一个功能强大的基于语义的图像检索系统

桂册 昨天 13:20

视频演示:

撸一个功能强大的基于语义的图像检索系统


 
大家好,这里是Coding茶水间。
在上一期视频中,我们介绍了Ultralytics框架下最新更新的语义检索功能,只需本地3行代码,就能构建一个基于文本语义的图像检索网站。
然而,框架自带的界面较为固定,如果需要自定义界面,还需额外开发;
此外,对中文支持不佳,使用中文检索时结果往往偏差;
还有小伙伴提出,能否实现以图搜图功能?
本期视频针对这些问题进行完善,如果你还没看过上一期,建议先回顾一下。
自定义界面的搭建

Ultralytics框架的VisualAISearch功能强大,但存在以下局限:

  • 界面固定:框架提供的界面无法满足个性化需求,需要自行编写GUI。
  • 中文支持差:直接用中文查询,语义理解不准,导致检索结果不理想。
  • 缺少以图搜图:原生仅支持文本检索,无法直接用图像查询相似图像。
本期我们构建了一个自定义界面,支持文本检索(兼容中文)和图像检索,全基于语义理解。
1.png

基于Kimi的中文转英文

我们使用PyQt5构建图形界面,集成Ultralytics的搜索功能,并借助Kimi AI处理中文翻译。
1. 中文测试

界面支持:

  • 文本检索:输入中文描述,如“一只欢乐的奔跑的狗”,系统后台自动翻译成英文后检索。
  • 图像检索:上传图像,系统基于语义相似度返回相似度最高的8个图像。
  • 显示相似度,并在网格中展示结果图像。
2.png

检索速度快,文本检索的主要耗时在翻译网络请求上,实际搜索实时完成。
2. 集成Kimi AI

上文中我们搜索的时候,直接输入的是中文。
但是原模型对中文支持比较差,我们用Kimi AI翻译输入为英文。新用户注册Kimi获15元额度,足够多次翻译。注册后创建API Key,填入代码即可。
3.png

翻译类(KimiTranslator.py):
python
  1. import requests
  2. class KimiTranslator:
  3.     def __init__(self, api_key):
  4.         self.api_key = api_key
  5.     def translate(self, text):
  6.         if not self.api_key:
  7.             raise ValueError("API key is empty")
  8.         
  9.         url = "https://api.moonshot.ai/v1/chat/completions"
  10.         headers = {
  11.             "Authorization": f"Bearer {self.api_key}",
  12.             "Content-Type": "application/json"
  13.         }
  14.         data = {
  15.             "model": "moonshot-v1-8k",
  16.             "messages": [
  17.                 {"role": "system", "content": "You are a helpful and precise assistant for translating Chinese to English."},
  18.                 {"role": "user", "content": f"Translate the following Chinese text to English: {text}"}
  19.             ],
  20.             "temperature": 0.3
  21.         }
  22.         
  23.         response = requests.post(url, headers=headers, json=data)
  24.         if response.status_code == 200:
  25.             return response.json()["choices"][0]["message"]["content"]
  26.         else:
  27.             raise Exception(f"Error: {response.status_code} - {response.text}")
复制代码
4.gif
核心代码实现

主窗口类(MainWindow)继承PyQt5,初始化搜索器并连接按钮信号。
完整代码:
python
  1. import sys
  2. from PyQt5 import QtWidgets
  3. import cv2
  4. from PyQt5.QtGui import *
  5. from PyQt5.QtCore import *
  6. import numpy as np
  7. from PyQt5.QtWidgets import *
  8. from ultralytics import solutions
  9. import os
  10. from main_window import Ui_MainWindow
  11. from KimiTranslator import KimiTranslator
  12. class MainWindow(QMainWindow, Ui_MainWindow):
  13.     def __init__(self):
  14.         super().__init__()
  15.         self.setupUi(self)
  16.         self.dir = "JPEGImages"
  17.         self.searcher = solutions.VisualAISearch(
  18.             data=self.dir
  19.         )
  20.         # 填写 Kimi 的 API Key
  21.         API_KEY = ""  # 在这里替换为你的 API Key
  22.         
  23.         # 创建翻译器实例
  24.         self.translator = KimiTranslator(API_KEY)
  25.         
  26.         # 初始化文件路径
  27.         self.file_path = None
  28.         
  29.         # 连接按钮信号(假设 UI 中的对象名;如果不同,请调整)
  30.         self.btnSelectImage.clicked.connect(self.selectImage)  # 选择图像按钮
  31.         self.btnTextSearch.clicked.connect(self.searchText2Pic)  # 以文搜图按钮
  32.         self.btnImageSearch.clicked.connect(self.searchPic2Pic)  # 以图搜图按钮
  33.     def updateImage(self, path, qlabel):
  34.         image = self.cv_imread(path)
  35.         if image is not None:
  36.             rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  37.             rows, cols, channels = rgb_image.shape
  38.             bytesPerLine = channels * cols
  39.             QImg = QImage(rgb_image.data, cols, rows, bytesPerLine, QImage.Format_RGB888)
  40.             qlabel.setPixmap(QPixmap.fromImage(QImg).scaled(qlabel.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
  41.     def selectImage(self):
  42.         file_path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Select an Image", "", "Image Files (*.jpg *.jpeg *.png)")
  43.         if file_path:
  44.             self.file_path = file_path
  45.             self.updateImage(self.file_path, self.queryImagePreview)
  46.     def cv_imread(self, path):
  47.         cv_img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
  48.         return cv_img
  49.     def searchText2Pic(self):
  50.         query_text = self.textQueryEdit.toPlainText().strip()
  51.         if not query_text:
  52.             QtWidgets.QMessageBox.warning(self, "提示", "输入文本不能为空")
  53.             return
  54.         translated = self.translator.translate(query_text)
  55.         results = self.searcher.search_text(translated)
  56.         self.updateResults(results)
  57.     def searchPic2Pic(self):
  58.         if not self.file_path:
  59.             QtWidgets.QMessageBox.warning(self, "提示", "请选择图像")
  60.             return
  61.         results = self.searcher.search_image(self.file_path)
  62.         self.updateResults(results)
  63.     def updateResults(self, results):
  64.         image_paths = []
  65.         similarities = []
  66.         for image_path, similarity in results[:8]:
  67.             if hasattr(image_path, '__str__'):
  68.                 image_path = str(image_path)
  69.             # 拼接完整路径
  70.             full_path = os.path.join(self.dir, image_path)
  71.             image_paths.append(full_path)
  72.             similarities.append(round(similarity, 4))
  73.         
  74.         # 更新图像和相似度,处理少于8个结果的情况
  75.         for i in range(8):
  76.             label = getattr(self, f'image_{i+1}')
  77.             sim_label = getattr(self, f'similarity_{i+1}')
  78.             if i < len(image_paths):
  79.                 self.updateImage(image_paths[i], label)
  80.                 sim_label.setText("相似度:" + str(similarities[i]))
  81.             else:
  82.                 label.clear()
  83.                 sim_label.setText("")
  84. if __name__ == "__main__":
  85.     app = QApplication(sys.argv)
  86.     window = MainWindow()
  87.     window.show()
  88.     sys.exit(app.exec_())
复制代码
5.gif
依赖:PyQt5、opencv-python、ultralytics、requests。图像目录“JPEGImages”需预置图像。
使用效果演示

文本检索示例

“一个人抱着一只狗”——包含人与狗的语义图像。
6.png

图像检索示例

上传一张实例图像,点击“以图搜图”,原图相似度1.0,其后是相似语义图像,如下图所示。
7.png

结语

这个系统完善了Ultralytics的语义检索,支持自定义界面和中文。
欢迎三连加关注,留言邮箱,我发源码给大家!如果有问题,评论区讨论。


来源:程序园用户自行投稿发布,如果侵权,请联系站长删除
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!

相关推荐

您需要登录后才可以回帖 登录 | 立即注册