撸一个功能强大的基于语义的图像检索系统
上一期视频中，我们介绍了Ultralytics框架下最新更新的语义检索功能，只需本地3行代码，就能构建一个基于文本语义的图像检索网站。
然而，框架自带的界面较为固定，如果需要自定义界面，还需额外开发；
此外，对中文支持不佳，使用中文检索时结果往往偏差；
还有小伙伴提出，能否实现以图搜图功能？
本期视频针对这些问题进行完善，如果你还没看过上一期，建议先回顾一下。
自定义界面的搭建
Ultralytics框架的VisualAISearch功能强大，但存在以下局限：
界面固定：框架提供的界面无法满足个性化需求，需要自行编写GUI。
中文支持差：直接用中文查询，语义理解不准，导致检索结果不理想。
缺少以图搜图：原生仅支持文本检索，无法直接用图像查询相似图像。
本期我们构建了一个自定义界面，支持文本检索（兼容中文）和图像检索，全基于语义理解。
主界面
基于Kimi的中文转英文
我们使用PyQt5构建图形界面，集成Ultralytics的搜索功能，并借助Kimi AI处理中文翻译。
1. 中文测试
界面支持：
文本检索：输入中文描述，如“一只欢乐的奔跑的狗”，系统后台自动翻译成英文后检索。
图像检索：上传图像，系统基于语义相似度返回相似度最高的8个图像。
显示相似度，并在网格中展示结果图像。
一只欢乐的奔跑的狗
检索速度快，文本检索的主要耗时在翻译网络请求上，实际搜索实时完成。
2. 集成Kimi AI
上文中我们搜索的时候，直接输入的是中文。
但是原模型对中文支持比较差，我们用Kimi AI翻译输入为英文。新用户注册Kimi获15元额度，足够多次翻译。注册后创建API Key，填入代码即可。
apikey
翻译类（KimiTranslator.py）：
python
import requests
class KimiTranslator:
def __init__(self, api_key):
self.api_key = api_key
def translate(self, text):
if not self.api_key:
raise ValueError("API key is empty")
url = "https://api.moonshot.ai/v1/chat/completions"
headers = {
"Authorization": f"Bearer {self.api_key}",
"Content-Type": "application/json"
}
data = {
"model": "moonshot-v1-8k",
"messages": [
{"role": "system", "content": "You are a helpful and precise assistant for translating Chinese to English."},
{"role": "user", "content": f"Translate the following Chinese text to English: {text}"}
],
"temperature": 0.3
}
response = requests.post(url, headers=headers, json=data)
if response.status_code == 200:
return response.json()["choices"][0]["message"]["content"]
else:
raise Exception(f"Error: {response.status_code} - {response.text}")
核心代码实现
主窗口类（MainWindow）继承PyQt5，初始化搜索器并连接按钮信号。
完整代码：
python
import sys
from PyQt5 import QtWidgets
import cv2
from PyQt5.QtGui import *
from PyQt5.QtCore import *
import numpy as np
from PyQt5.QtWidgets import *
from ultralytics import solutions
import os
from main_window import Ui_MainWindow
from KimiTranslator import KimiTranslator
class MainWindow(QMainWindow, Ui_MainWindow):
def __init__(self):
super().__init__()
self.setupUi(self)
self.dir = "JPEGImages"
self.searcher = solutions.VisualAISearch(
data=self.dir
)
# 填写 Kimi 的 API Key
API_KEY = "" # 在这里替换为你的 API Key
# 创建翻译器实例
self.translator = KimiTranslator(API_KEY)
# 初始化文件路径
self.file_path = None
# 连接按钮信号（假设 UI 中的对象名；如果不同，请调整）
self.btnSelectImage.clicked.connect(self.selectImage) # 选择图像按钮
self.btnTextSearch.clicked.connect(self.searchText2Pic) # 以文搜图按钮
self.btnImageSearch.clicked.connect(self.searchPic2Pic) # 以图搜图按钮
def updateImage(self, path, qlabel):
image = self.cv_imread(path)
if image is not None:
rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
rows, cols, channels = rgb_image.shape
bytesPerLine = channels * cols
QImg = QImage(rgb_image.data, cols, rows, bytesPerLine, QImage.Format_RGB888)
qlabel.setPixmap(QPixmap.fromImage(QImg).scaled(qlabel.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))
def selectImage(self):
file_path, _ = QtWidgets.QFileDialog.getOpenFileName(self, "Select an Image", "", "Image Files (*.jpg *.jpeg *.png)")
if file_path:
self.file_path = file_path
self.updateImage(self.file_path, self.queryImagePreview)
def cv_imread(self, path):
cv_img = cv2.imdecode(np.fromfile(path, dtype=np.uint8), -1)
return cv_img
def searchText2Pic(self):
query_text = self.textQueryEdit.toPlainText().strip()
if not query_text:
QtWidgets.QMessageBox.warning(self, "提示", "输入文本不能为空")
return
translated = self.translator.translate(query_text)
results = self.searcher.search_text(translated)
self.updateResults(results)
def searchPic2Pic(self):
if not self.file_path:
QtWidgets.QMessageBox.warning(self, "提示", "请选择图像")
return
results = self.searcher.search_image(self.file_path)
self.updateResults(results)
def updateResults(self, results):
image_paths = []
similarities = []
for image_path, similarity in results[:8]:
if hasattr(image_path, '__str__'):
image_path = str(image_path)
# 拼接完整路径
full_path = os.path.join(self.dir, image_path)
image_paths.append(full_path)
similarities.append(round(similarity, 4))
# 更新图像和相似度，处理少于8个结果的情况
for i in range(8):
label = getattr(self, f'image_{i+1}')
sim_label = getattr(self, f'similarity_{i+1}')
if i < len(image_paths):
self.updateImage(image_paths[i], label)
sim_label.setText("相似度：" + str(similarities[i]))
else:
label.clear()
sim_label.setText("")
if __name__ == "__main__":
app = QApplication(sys.argv)
window = MainWindow()
window.show()
sys.exit(app.exec_())
依赖：PyQt5、opencv-python、ultralytics、requests。图像目录“JPEGImages”需预置图像。
使用效果演示
文本检索示例
“一个人抱着一只狗”——包含人与狗的语义图像。
一个人抱着一只狗