手把手教你搭建一个汽车型号图片搜索系统

前言

该项目将分两个阶段执行。首先，我们将收集数据并将其转换为矢量。在第二阶段，我们将使用该数据以及输入图像，使用Streamlit框架显示类似的图像。

配置环境

我们将使用ImageBind，这是 Meta 开发的开源库，用于将图像转换为矢量。

代码首先安装 ImageBind 库，该库需要从其 GitHub 存储库克隆才能正确集成和使用。

git clone https://github.com/facebookresearch/ImageBind.git cd ImageBind pip install -e 。

此外，我们还需要一些其他库，包括ultralytics和qdrant-client，以确保项目正确高效地运行。

pip install ultralyticspip install qdrant-clientpip install streamlit

数据采集

我们收集了一系列代表各种类型汽车及其各自价格的图片。

请从这里访问数据集：https://www.kaggle.com/datasets/vanshkhaneja/cars-prices/

随后，将它们存储在列表中：一个用于存储名称，另一个用于存储各自的价格。

cars_img_list = [ "img01" , "img02" , "img03" , "img04" , "img05" , "img06" , "img07" , "img08" , "img09" , "img10" , "img11" , "img12" , "img13" , "img14" , "img15" ] cars_cost_list = [ "6.49" , "3.99" , "6.66" , "6.65" , "7.04" , "5.65" , "61.85" , "11.00" , "11.63" , "11.56" , "11.86" , "46.05" , "75.90" , “13.59”，“13.99” ]

导入库

现在让我们导入将图像转换为嵌入所需的所有必要库。

from ultralytics import YOLOimport cv2import os
import torchfrom imagebind import datafrom imagebind.models import imagebind_modelfrom imagebind.models.imagebind_model import ModalityType

图像分割

我们将使用YOLOv8算法裁剪掉汽车，从而从图像中去除不必要的噪音。

为了实现这一点，我们首先借助YOLOv8在汽车周围绘制一个边界框，然后使用OpenCV剪切出该区域，最后将图像保存在目录中。

model = YOLO('yolov8n.pt')
for im in cars_img_list:    img = cv2.imread("cars_imgs/"+im+".jpg")    img = cv2.resize(img,(320,245))    results = model(img,stream=True)
    for r in results:        for box in r.boxes:            x1,y1,x2,y2 = box.xyxy[0]            x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
            cv2.rectangle(img,(int(x1),int(y1)),(int(x2),int(y2)),(255,0,0),1)        cv2.imwrite("cropped_imgs/"+im+"_cropped.jpg",img[y1:y2, x1:x2])

图像转矢量

接下来，我们将使用ImageBind库将裁剪的图像转换为矢量嵌入，从而将其转换为数字格式。

注意：这是一个耗时的过程，因为它将从互联网上下载模型。

embedding_list = [] 
model_embed = imagebind_model.imagebind_huge(pretrained= True ) model_embed.eval ( ) model_embed.to( "cpu" ) 
for i in  range ( 1 , len (cars_img_list)):     img_path = "cropped_imgs/img" + str (i)+ "_cropped.jpg"     vision_data = data.load_and_transform_vision_data([img_path], device) 
    with torch.no_grad():         image_embeddings = model_embed({ModalityType.VISION: vision_data})     embedding_list.append(image_embeddings)

为了减少处理时间，您可以设置pretrained=False，尽管这会降低模型的准确性。

让我们保存嵌入以供稍后在代码中使用。

import picklewith open('embedded_data.pickle', 'wb') as file:    pickle.dump(embedding_list, file)

相似图片搜索

进入代码的第二阶段，我们将继续以图像作为输入并识别最相似的图像。随后，我们将显示这些图像及其各自的价格。

导入库

除了先前导入的库之外，还需要以下附加库：

import streamlit as stfrom PILimport Imageimport base64import osfrom io import BytesIOimport numpy as np

现在我们将通过启动ImageBind模型来启动代码，该模型会将上传的输入图像转换为矢量嵌入。

model_embed = imagebind_model.imagebind_huge(pretrained=True)model_embed.eval()model_embed.to("cpu")

让我们继续打开保存的文件“embedded_data.pickle”，其中包含我们的图像数据集的矢量数据。

with open('embedded_data.pickle', 'rb') as file:    embedding_list = pickle.load(file)

存储矢量数据

我们将利用开源矢量数据库Qdrant来存储所有图像的嵌入并将其与输入图像进行比较。

client = QdrantClient(":memory:")
client.recreate_collection(    collection_name='vector_comparison',    vectors_config=VectorParams(size=1024, distance=Distance.COSINE))
client.upsert(    collection_name='vector_comparison',    points=[        PointStruct(id=i, vector=embedding_list[i]['vision'][0].tolist()) for i in range(15)    ])

比较图像

接下来，我们将把存储在Qdrant数据库中的每个向量嵌入与提供给程序的输入图像进行比较。

这将通过 3 个步骤完成：

在这个过程中，我们采用余弦相似度来评估嵌入之间的相似度。我们声明了一个函数，该函数以图像作为输入，并给出与输入图像最相似的 4 张图像的索引。

def image_to_similar_index(cv2Image):        img = cv2.resize(cv2Image,(320,245))        model = YOLO('yolov8n.pt')        results = model(img,stream=True)        results = model(img,stream=True)for r in results:for box in r.boxes:                    x1,y1,x2,y2 = box.xyxy[0]                    x1,y1,x2,y2 = int(x1),int(y1),int(x2),int(y2)
                    cv2.rectangle(img,(int(x1),int(y1)),(int(x2),int(y2)),(255,0,0),1)                    cropped_img = img[y1:y2, x1:x2]                cv2.imwrite("test_cropped.jpg",cropped_img)
        vision_data = data.load_and_transform_vision_data(["test_cropped.jpg"], device)        with torch.no_grad():            test_embeddings = model_embed({ModalityType.VISION: vision_data})
        client.upsert(                collection_name='vector_comparison',                points=[                  PointStruct(id=20, vector=test_embeddings['vision'][0].tolist()),                ])        search_result = client.search(            collection_name='vector_comparison',            query_vector=test_embeddings['vision'][0].tolist(),            limit=20# Retrieve top similar vectors (excluding the new vector itself)        )return [search_result[1].id,search_result[2].id,search_result[3].id,search_result[4].id]

部署模型

我们现在将继续为我们的模型开发前端 Web 应用程序，以增强交互性和用户友好性。为了实现这一点，我们将利用Streamlit以简单有效的方式为我们的 Python 应用程序创建 Web 界面。我们将首先配置页面并将文件上传器小部件集成到网页上。

st.set_page_config(layout="wide")st.title('Similar Cars Finder')st.markdown("""<style>.block-container {padding-top: 3rem;padding-bottom: 0rem;padding-left: 5rem;padding-right: 5rem;                }</style>        """, unsafe_allow_html=True)uploaded_file = st.file_uploader("Upload an image of a car", type=["jpg", "jpeg", "png"])

现在我们将创建一个函数来显示带有适当填充和边距的图像以及价格。此函数将图像和价格表作为输入，并以格式化的方式将其显示在网页上。

defdisplay_images_with_padding_and_price(images, prices, width, padding, gap):    cols = st.columns(len(images))for col, img, price in zip(cols, images, prices):with col:            col.markdown(f"""                <div style="margin-right: {0}px; text-align: center;">                    <img src="data:image/jpeg;base64,{img}" width="{250}px;margin-right: {50}px; ">                    <p style="font-size: 20px;">₹{price} Lakhs</p>                </div>                 """,                unsafe_allow_html=True,            )

最后，我们将读取上传的图像作为输入，将其转换为NumPy数组，并将其作为输入提供给image_to_similar_index我们之前定义的函数，该函数将返回与输入最相似的图像的索引。然后，我们将检索与返回的索引相对应的图像和价格，并将它们提供给函数display_images_with_padding_and_price，该函数将格式化图像并将其显示在网页上。

if uploaded_file isnotNone:    car_image = Image.open(uploaded_file)    img_array = np.array(car_image)    st.image(car_image, caption='Uploaded Car Image', use_column_width=False, width=300)    results = image_to_similar_index(img_array)
if os.path.exists("cars_imgs"):        car_images = [os.path.join(car_images_dir, img) for img in os.listdir(car_images_dir) if img.endswith(('jpg', 'jpeg', 'png'))]        print(car_images)else:        st.error(f"Directory {car_images_dir} does not exist")        car_images = []
if len(car_images) < 4:        st.error("Not enough car images in the local storage")else:        car_imagess = []for i in results:             car_imagess.append(car_images[i])        car_prices = [cars_cost_list[a] for a in results]
        car_images_pil = []for img_path in car_imagess:try:                img = Image.open(img_path)                buffered = BytesIO()                img.save(buffered, format="JPEG")                img_str = base64.b64encode(buffered.getvalue()).decode()                car_images_pil.append(img_str)except Exception as e:                st.error(f"Error processing image {img_path}: {e}")
if car_images_pil:            st.subheader('Similar Cars with Prices')            display_images_with_padding_and_price(car_images_pil, car_prices, width=200, padding=10, gap=20)

最终输出

将图片上传到网页后，ImageBind模型会开始加载，这可能需要一点时间。但是，一旦模型完全加载，图片就会转换为嵌入，并与其他图片进行比较，以找出最相似的图片。随后，相似的图片就会显示在网页上。

完整代码：

https://github.com/vansh-khaneja/Car-Similarity-Search