当前位置： 首页 > article >正文

resnet50,clip，Faiss+Flask简易图文搜索服务

article 2024/11/23 14:57:09

一、实现

文件夹目录结构：

templates

        -----upload.html

faiss_app.py

前端代码：

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Search and Show Multiple Images</title>
    <style>
        #image-container {
            display: flex;
            flex-wrap: wrap;
        }
        #image-container img {
            max-width: 150px;
            margin: 10px;
        }
    </style>
</head>
<body>
    <h1>Search Images</h1>
     <!-- 上传表单 -->
     <form id="upload-form" enctype="multipart/form-data">
        <input type="file" id="file-input" name="file" accept="image/*" required>
        <input type="submit" value="Upload">
    </form>

    <!-- 搜索框 -->
    <form id="search-form">
        <input type="text" id="search-input" name="query" placeholder="Enter search term" required>
        <input type="submit" value="Search">
    </form>

    <h2>Search Results</h2>
    <!-- 显示搜索返回的多张图片 -->
    <div id="image-container"></div>

    <!-- 使用JS处理表单提交 -->
    <script>
        document.getElementById('search-form').addEventListener('submit', async function(event) {
            event.preventDefault();  // 阻止表单默认提交行为
            
            const query = document.getElementById('search-input').value;  // 获取搜索框中的输入内容

            try {
                // 发送GET请求，将搜索关键词发送到后端
                const response = await fetch(`/search?query=${encodeURIComponent(query)}`, {
                    method: 'GET',
                });

                // 确保服务器返回JSON数据
                const data = await response.json();

                // 清空图片容器
                const imageContainer = document.getElementById('image-container');
                imageContainer.innerHTML = '';

                // 遍历后端返回的图片URL数组，动态创建<img>标签并渲染
                data.image_urls.forEach(url => {
                    const imgElement = document.createElement('img');
                    imgElement.src = url;  // 设置图片的src属性为返回的URL
                    imageContainer.appendChild(imgElement);  // 将图片添加到容器中
                });
            } catch (error) {
                console.error('Error searching for images:', error);
            }
        });
        document.getElementById('upload-form').addEventListener('submit', async function(event) {
            event.preventDefault();  // 阻止表单默认提交行为
            
            const fileInput = document.getElementById('file-input');
            const formData = new FormData();
            formData.append('file', fileInput.files[0]);  // 获取用户上传的图片文件

            try {
                // 发送POST请求，将图片发送到后端
                const response = await fetch('/search_by_images', {
                    method: 'POST',
                    body: formData
                });

                // 确保服务器返回JSON数据
                const data = await response.json();

                // 清空图片容器
                const imageContainer = document.getElementById('image-container');
                imageContainer.innerHTML = '';

                // 遍历后端返回的图片URL数组，动态创建<img>标签并渲染
                data.image_urls.forEach(url => {
                    const imgElement = document.createElement('img');
                    imgElement.src = url;  // 设置图片的src属性为返回的URL
                    imageContainer.appendChild(imgElement);  // 将图片添加到容器中
                });
            } catch (error) {
                console.error('Error uploading file:', error);
            }
        });
    </script>
</body>
</html>

后端代码：

from sentence_transformers import SentenceTransformer, util
from torchvision import models, transforms
from PIL import Image
from flask import Flask, request, jsonify, current_app, render_template, send_from_directory, url_for
from werkzeug.utils import secure_filename
import faiss
import os, glob
import numpy as np
from markupsafe import escape
import shutil

#Load CLIP model
model = SentenceTransformer('clip-ViT-B-32')
IMAGE_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.gif', '.bmp'}

UPLOAD_FOLDER = 'uploads/'
IMAGES_PATH  = "C:\\Users\\cccc\\Pictures\\cls_auto_config"

def generate_clip_embeddings(images_path, model):
    image_paths = []
    # 使用 os.walk 遍历所有子目录和文件
    for root, dirs, files in os.walk(images_path):
        for file in files:
            # 获取文件的扩展名并转换为小写
            ext = os.path.splitext(file)[1].lower()
            # 判断是否是图片文件
            if ext in IMAGE_EXTENSIONS:
                image_paths.append(os.path.join(root, file)) 
    embeddings = []
    for img_path in image_paths:
        image = Image.open(img_path)
        embedding = model.encode(image)
        embeddings.append(embedding)
    
    return embeddings, image_paths

def generate_res50_embeddings(images_path):
    # Load the pretrained model
    res50_model = models.resnet50(pretrained=True)
    res50_model = res50_model.eval()

    # Define the image transformations
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    image_paths = []
    # 使用 os.walk 遍历所有子目录和文件
    for root, dirs, files in os.walk(images_path):
        for file in files:
            # 获取文件的扩展名并转换为小写
            ext = os.path.splitext(file)[1].lower()
            # 判断是否是图片文件
            if ext in IMAGE_EXTENSIONS:
                image_paths.append(os.path.join(root, file)) 
    embeddings = []
    for img_path in image_paths:
        image = Image.open(img_path)
        # Apply the transformations and get the image vector
        image = transform(image).unsqueeze(0)
        image_vector = res50_model(image).detach().numpy()
        embeddings.append(image_vector[0])
    return embeddings, image_paths

def create_faiss_index(embeddings, image_paths, output_path):

    dimension = len(embeddings[0])

    # 分情况创建Faiss索引对象
    if len(image_paths) < 39 * 256:
        # 如果条目很少，直接用最普通的L2索引
        faiss_index = faiss.IndexFlatL2(dimension)
    elif len(image_paths) < 39 * 4096:
        # 如果条目少于39 × 4096，就只用PQ量化，不使用IVF
        faiss_index = faiss.index_factory(dimension, 'OPQ64_256,PQ64x8')
    else:
        # 否则就加上IVF
        faiss_index = faiss.index_factory(dimension, 'OPQ64_256,IVF4096,PQ64x8')
    res = faiss.StandardGpuResources()
    co = faiss.GpuClonerOptions()
    co.useFloat16 = True
    faiss_index = faiss.index_cpu_to_gpu(res, 0, faiss_index, co)

    #index = faiss.IndexFlatIP(dimension)
    faiss_index = faiss.IndexIDMap(faiss_index)
    
    vectors = np.array(embeddings).astype(np.float32)

    # Add vectors to the index with IDs
    faiss_index.add_with_ids(vectors, np.array(range(len(embeddings))))
    
    # Save the index
    faiss_index = faiss.index_gpu_to_cpu(faiss_index)
    faiss.write_index(faiss_index, output_path)
    print(f"Index created and saved to {output_path}")
    
    # Save image paths
    with open(output_path + '.paths', 'w') as f:
        for img_path in image_paths:
            f.write(img_path + '\n')
    
    return faiss_index

def load_faiss_index(index_path):
    faiss_index = faiss.read_index(index_path)
    with open(index_path + '.paths', 'r') as f:
        image_paths = [line.strip() for line in f]
    print(f"Index loaded from {index_path}")
    if not faiss_index.is_trained:
            raise RuntimeError(f'从[{index_path}]加载的Faiss索引未训练')
    res = faiss.StandardGpuResources()
    co = faiss.GpuClonerOptions()
    co.useFloat16 = True
    faiss_index = faiss.index_cpu_to_gpu(res, 0, faiss_index, co)
    return faiss_index, image_paths


def retrieve_similar_images(query, model, index, image_paths, top_k=3):
    
    # query preprocess:
    if query.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
        query = Image.open(query)

    query_features = model.encode(query)
    query_features = query_features.astype(np.float32).reshape(1, -1)

    distances, indices = index.search(query_features, top_k)

    retrieved_images = [image_paths[int(idx)] for idx in indices[0]]

    return query, retrieved_images

def retrieve_res50_similar_images(query, index, image_paths, top_k=3):
    # query preprocess:
    if query.endswith(('.png', '.jpg', '.jpeg', '.tiff', '.bmp', '.gif')):
        image = Image.open(query)
        # Load the pretrained model
        res50_model = models.resnet50(pretrained=True)
        res50_model = res50_model.eval()

        # Define the image transformations
        transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
        ])

        # Apply the transformations and get the image vector
        image = transform(image).unsqueeze(0)
        query_features = res50_model(image).detach().numpy()
        query_features = query_features[0]
        query_features = query_features.astype(np.float32).reshape(1, -1)

        distances, indices = index.search(query_features, top_k)

        retrieved_images = [image_paths[int(idx)] for idx in indices[0]]

        return query, retrieved_images

# 检查文件扩展名是否允许
def allowed_file(filename):
    return '.' in filename and "." + filename.rsplit('.', 1)[1].lower() in IMAGE_EXTENSIONS

def search():
    query = request.args.get('query')  # 获取搜索关键词
    safe_query = escape(query)

    if not query:
        return jsonify({"error": "No search query provided"}), 400
    index, image_paths = None, []
    OUTPUT_INDEX_PATH = f"{app.config['UPLOAD_FOLDER']}/vector.index"
    if os.path.exists(OUTPUT_INDEX_PATH):
        index, image_paths = load_faiss_index(OUTPUT_INDEX_PATH)
    else:
        # embeddings, image_paths = generate_clip_embeddings(IMAGES_PATH, model)
        embeddings, image_paths = generate_res50_embeddings(IMAGES_PATH)
        index = create_faiss_index(embeddings, image_paths, OUTPUT_INDEX_PATH)
    query, retrieved_images = retrieve_similar_images(query, model, index, image_paths, top_k=5)


    image_urls = []
    for path in retrieved_images:
        base_name = os.path.basename(path)
        shutil.copy(path, os.path.join(app.config['UPLOAD_FOLDER'], base_name))
        image_urls.append(url_for('uploaded_file_path', filename=base_name))

    return jsonify({"image_urls": image_urls})


def search_by_images():
    # 检查请求中是否有文件
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400
    file = request.files['file']

    # 检查文件是否为空
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400
    print(file.filename)
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)

        index, image_paths = None, []
        OUTPUT_INDEX_PATH = f"{app.config['UPLOAD_FOLDER']}/images_vector.index"
        if os.path.exists(OUTPUT_INDEX_PATH):
            index, image_paths = load_faiss_index(OUTPUT_INDEX_PATH)
        else:
            embeddings, image_paths = generate_res50_embeddings(IMAGES_PATH)
            index = create_faiss_index(embeddings, image_paths, OUTPUT_INDEX_PATH)
        filepath, retrieved_images = retrieve_res50_similar_images(filepath, index, image_paths, top_k=5)

        image_urls = []
        for path in retrieved_images:
            base_name = os.path.basename(path)
            shutil.copy(path, os.path.join(app.config['UPLOAD_FOLDER'], base_name))
            image_urls.append(url_for('uploaded_file_path', filename=base_name))

        return jsonify({"image_urls": image_urls})
    else:
        return jsonify({"error": "Invalid file"}), 400


def index():
    return render_template('upload.html')

# 提供静态文件的访问路径
def uploaded_file_path(filename):
    return send_from_directory(app.config['UPLOAD_FOLDER'], filename)

if __name__ == "__main__":
    app = Flask(__name__)
    app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)
    # 主页显示上传表单
    app.route('/')(index)
    app.route('/search', methods=['GET'])(search)
    app.route('/uploads/images/<filename>')(uploaded_file_path)
    app.route('/search_by_images', methods=['POST'])(search_by_images)
    app.run(host='0.0.0.0', port=8080, debug=True)

二、实现效果

三、参考文章

1. https://towardsdatascience.com/building-an-image-similarity-search-engine-with-faiss-and-clip-2211126d08fa

2.向量数据库Faiss的搭建与使用 - 很久8899 - 博客园


http://www.kler.cn/a/405877.html

相关文章：

  • FPGA实现串口升级及MultiBoot（九)BPI FLASH相关实例演示
  • 小程序24-滚动效果：scroll-view组件详解
  • 【企业级分布式系统】ZooKeeper集群
  • 杰理-gpadc
  • 一次封装，解放双手：Requests如何实现0入侵请求与响应的智能加解密
  • ss 命令的基本用法
  • 信息收集系列（六）：路径爬取与目录爆破
  • Flutter踩坑记录（三）-- 更改入口执行文件
  • 【MySQL实战45讲笔记】基础篇——事务隔离
  • OCR的技术发展及OCR厂商
  • E - 11/22 Subsequence题解
  • nvm安装node遇到的若干问题（vscode找不到npm文件、环境变量配置混乱、npm安装包到D盘）
  • 图像预处理之图像滤波
  • React-useEffect的使用
  • 任务通知的本质(任务通知车辆运行) 软件定时器的本质（增加游戏音效）
  • MybatisPlus编写join查询
  • 记录下jekins新建个前端部署配置项
  • 单片机学习笔记 9. 8×8LED点阵屏
  • ffmpeg区域颜色覆盖
  • Echarts中柱状图完成横向布局
  • 【MySQL】数据库的隔离级
  • 使用 vscode 调试 nodejs 代码
  • 数据分析-51-时间序列分解之局部均值分解LMD
  • 【Three.js基础学习】28.Coffee Smoke
  • 鸿蒙进阶篇-TextInputTextArea和Checkbox
  • E. Counting Arrays