上传 检索 接口

app/api/routes.py

 from flask import jsonify, request
 from app.api import api_bp
+from app.services.feature_extractor import FeatureExtractor
+from app.services.image_search import ImageSearch
+from app.services.image_upload import ImageUpload
+from app.services.milvus import MilvusClient
+
 
 @api_bp.route('/hello', methods=['GET'])
 def hello_world():
@@ -24,9 +29,29 @@ def upload():
     # 访问具体字段
     bucket = data.get('bucket')
     image2keys = data.get('image2keys')
-    print(bucket)
-    print(image2keys)
+
+    feature_extractor = FeatureExtractor()
+    milvus = MilvusClient().connect()
+    image_upload = ImageUpload(feature_extractor, milvus)
+    image_upload.upload_many(bucket, image2keys)
     return jsonify({
-        'message': f'Hello!',
-        'status': 'success'
+        'code': 0,
+        'message': '',
+    })
+
+@api_bp.route('/search', methods=['GET'])
+def search():
+    image = request.args.get("image")
+    top_k = request.args.get("top_k", type=int)
+    bucket = request.args.get("bucket")
+
+    feature_extractor = FeatureExtractor()
+    milvus = MilvusClient().connect()
+
+    result = ImageSearch(feature_extractor, milvus).image_to_image_search(bucket, image, top_k)
+
+    return jsonify({
+        'code': 0,
+        'message': '',
+        'data': result,
     })

app/services/image_search.py

@@ -10,7 +10,7 @@ class ImageSearch:
     def image_to_image_search(self, bucket, image, top_k = 100):
         try:
             # 提取查询图像的特征
-            vector = self.feature_extractor.extract_features(image)
+            vector = self.feature_extractor.extract_from_url(image)
 
             results = self.milvus.search(bucket, vector, top_k)
 
@@ -20,9 +20,10 @@ class ImageSearch:
 
             # 返回结果
             keys = [hit.entity.get("key") for hit in results[0]]
+            images = [hit.entity.get("image") for hit in results[0]]
             scores = [hit.score for hit in results[0]]
 
-            return list(zip(keys, scores))
+            return list(zip(keys, images, scores))
 
         except Exception as e:
             self.__logger.error(f"Image search failed: {e}")

app/services/image_upload.py

 import logging
 
-class Upload:
+class ImageUpload:
     __logger = logging.getLogger(__name__)
 
     def __init__(self, feature_extractor, milvus):
@@ -10,12 +10,14 @@ class Upload:
     def upload_one(self, bucket, image, key):
         self.upload_many(bucket, {image: key})
 
-    def upload_many(self, bucket, image2key):
+    def upload_many(self, bucket, image2keys):
         images = []
         keys = []
         vectors = []
 
-        for image, key in image2key.items():
+        for image2key in image2keys:
+            image = image2key["image"]
+            key = image2key["key"]
             vector = self.feature_extractor.extract_from_url(image)
             images.append(image)
             keys.append(key)