milvus+flask山寨《从零构建向量数据库》第7章case1

常规练手，图片搜刮山寨版。拜读罗云大佬著作，结果只有操纵层的东西可以上上手。
书中是自己写的向量数据库，这边直接用python拼个现成的milvus向量数据库。
1. 创建一个向量数据库以及对应的相应数据表：

1. # Milvus Setup Arguments
2. COLLECTION_NAME = 'animal_search'
3. DIMENSION = 2048
4. MILVUS_HOST = "localhost"
5. MILVUS_PORT = "19530"
7. # Inference Arguments
8. BATCH_SIZE = 128
10. from pymilvus import connections
12. # Connect to the instance
13. connections.connect(host=MILVUS_HOST,port=MILVUS_PORT)
15. from pymilvus import utility
17. # Remove any previous collection with the same name
18. if utility.has_collection(COLLECTION_NAME):
19.     utility.drop_collection(COLLECTION_NAME)
21. #创建保存ID、图片文件路径及Embeddings的Collection。
22. from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
24. fields = [
25.         FieldSchema(name='id',dtype=DataType.INT64, is_primary=True, auto_id=True),
26.         FieldSchema(name='filepath', dtype=DataType.VARCHAR,max_length=200),
27.         FieldSchema(name='image_embedding',dtype=DataType.FLOAT_VECTOR,dim=DIMENSION)
28.         ]
29. schema = CollectionSchema(fields=fields)
30. collection = Collection(name=COLLECTION_NAME, schema=schema)
32. index_params = {
33.         'metric_type':'L2',
34.         'index_type': "IVF_FLAT",
35.         'params':{'nlist':16384}
36. }
37. collection.create_index(field_name="image_embedding",index_params=index_params)
38. collection.load()

复制代码

2. 写一堆图片进去存着，向量其实就是各种像素间的维度特征，

1. # Milvus Setup Arguments
2. COLLECTION_NAME = 'animal_search'
3. DIMENSION = 2048
4. MILVUS_HOST = "localhost"
5. MILVUS_PORT = "19530"
7. # Inference Arguments
8. BATCH_SIZE = 128
10. from pymilvus import connections
12. # Connect to the instance
13. connections.connect(host=MILVUS_HOST, port=MILVUS_PORT)
15. import glob
17. paths = glob.glob('/mcm/vectorDB_training/animals_db/*',recursive=True)
19. #分批预处理数据
20. import torch
21. # Load the embedding model with the last layer removed
22. model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet50', pretrained=True)
23. model = torch.nn.Sequential(*(list(model.children())[:-1]))
24. model.eval()
26. from torchvision import transforms
27. # Preprocessing for images
28. preprocess = transforms.Compose([
29.     transforms.Resize(256),
30.     transforms.CenterCrop(224),
31.     transforms.ToTensor(),
32.     transforms.Normalize(mean=[0.485,0.456,0.406],std=[0.229,0.224,0.225])
33. ])
35. #插入数据
36. from PIL import Image
37. from tqdm import tqdm
39. # Embed function that embeds the batch and inserts it
40. def embed(data):
41.     from pymilvus import FieldSchema, CollectionSchema, DataType, Collection
43.     fields = [
44.         FieldSchema(name='id',dtype=DataType.INT64, is_primary=True, auto_id=True),
45.         FieldSchema(name='filepath', dtype=DataType.VARCHAR,max_length=200),
46.         FieldSchema(name='image_embedding',dtype=DataType.FLOAT_VECTOR,dim=DIMENSION)
47.         ]
48.     schema = CollectionSchema(fields=fields)
49.     collection = Collection(name=COLLECTION_NAME, schema=schema)
50.     with torch.no_grad():
51.         output = model(torch.stack(data[0])).squeeze()
52.         collection.insert([data[1],output.tolist()])
53.     collection.flush()
55. data_batch = [[],[]]
57. # Read the images into batches for embedding and insertion
58. for path in tqdm(paths):
59.     im = Image.open(path).convert('RGB')
60.     data_batch[0].append(preprocess(im))
61.     data_batch[1].append(path)
62.     if len(data_batch[0]) % BATCH_SIZE == 0:
63.         embed(data_batch)
64.         data_batch = [[],[]]
66. # Embed and insert the remainder
67. if len(data_batch[0]) != 0:
68.     embed(data_batch)

复制代码

3. 向量化图片的函数要单独拎出来，做搜刮功能的时候用它。

1. import torch
2. import torchvision.transforms as transforms
3. from torchvision.models import resnet50
4. from PIL import Image
6. def extract_features(image_path):
7.     # 加载预训练的 ResNet-50 模型
8.     model = resnet50(pretrained=True)
9.     model = torch.nn.Sequential(*list(model.children())[:-1])  #移除fc层，不移除，向量最后就是1000层，而不是2048
10.     model.eval()
12.     # 图像预处理
13.     preprocess = transforms.Compose([
14.         transforms.Resize(256),
15.         transforms.CenterCrop(224),
16.         transforms.ToTensor(),
17.         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
18.     ])
20.     # 读取图像
21.     img = Image.open(image_path)
22.     img_t = preprocess(img)
23.     batch_t = torch.unsqueeze(img_t, 0)
25.     # 提取特征
26.     with torch.no_grad():
27.         out = model(batch_t)
29.     # 将特征向量转换为一维数组并返回
30.     return out.flatten().numpy()

复制代码

4. 用flask做的界面

1. from flask import Flask,request,jsonify
2. from flask import render_template
3. from image_eb import extract_features
4. #from pymilvus import connections
5. from pymilvus import MilvusClient
7. import logging
8. import os
9. import shutil
11. MILVUS_HOST = "localhost"
12. MILVUS_PORT = "19530"
13. COLLECTION_NAME = 'animal_search'
14. TOP_K = 3
16. app = Flask(__name__)
17. milvus_client = MilvusClient(uri="http://localhost:19530")
19. @app.route("/")
20. def index():
21.     return render_template("index.html")
23. @app.route("/upload",methods=["POST"])
24. def upload_image():
25.     image_file = request.files["image"]
26.     image_id_str = request.form.get("image_id")
27.     data = []
28.     #检查image_id是否存在。
29.     if not image_id_str:
30.         return jsonify({"message": "Image ID is required"}),400
31.     #image id转化为整型
32.     try:
33.         image_id = int(image_id_str)
34.         data.append(image_id)
35.     except ValueError:
36.         return jsonify({"message": "Invalid image ID. It must be an integer"}),400
37.     filename = image_file.filename
38.     image_path = os.path.join("static/images",image_id_str)
39.     image_file.save(image_path)
40.     image_features = extract_features(image_path)
41.     data.append(image_features)
42.     data_dict = dict(filepath=image_path,image_embedding=data[1])
43.     #更新数据库中记录
44.     milvus_client.insert(collection_name=COLLECTION_NAME,data=[data_dict])
45.     return jsonify({"message": "Image uploaded successfully", "id": image_id})
47. @app.route("/search",methods=["POST"])
48. def search_image():
49.     image_file = request.files["image"]
50.     image_path = os.path.join("static/images","temp_image.jpg")
51.     image_file.save(image_path)
52.     image_features = extract_features(image_path)
53.     data_li = [extract_features(image_path).tolist()]
54.     search_result = milvus_client.search(
55.         collection_name=COLLECTION_NAME,
56.         data=data_li,
57.         output_fields=["filepath"],
58.         limit=TOP_K,
59.         search_params={'metric_type': 'L2', 'params': {}},
60.     )
61.     dict_search_result = search_result[0]
62.     arr_search_result = []
63.     destination_folder = '/mcm/vectorDB_training/static/images'
64.     for index,value in enumerate(dict_search_result):
65.         source_file = value["entity"]["filepath"]
66.         base_file_name = os.path.basename(source_file)
67.         destination_file = os.path.join(destination_folder, base_file_name)
68.         shutil.copy(source_file, destination_file)
69.         key_file_name = os.path.join("/static/images",base_file_name)
71.         arr_search_result.append(key_file_name)        
73.     image_urls = [
74.             f"{filepath}" for filepath in arr_search_result
75.         ]
76.     return jsonify({"image_urls":image_urls})
78. if __name__=="__main__":
79.     app.run(host='0.0.0.0',port=5020,debug=True)

复制代码

谁私信我说脸谱不是猫，只是缩略图里没放清楚，睁大go眼瞧瞧：

小网站结构，以及其他杂代码，可以查看以及直接下载：https://www.ituring.com.cn/book/3305

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。