相似匹配接口

app/service/similar_match/service.py

@@ -10,17 +10,21 @@ from pymilvus import MilvusClient
 from app.core.config import *
 from torchvision import transforms
 
+from app.schemas.similar_match import SimilarMatchMItem
 from app.service.utils.decorator import RunTime
 
 
 class SimilarMatch:
-    def __init__(self):
+    def __init__(self, request_data):
         self.minio_client = Minio(
             f"{MINIO_IP}:{MINIO_PORT}",
             access_key=MINIO_ACCESS,
             secret_key=MINIO_SECRET,
             secure=MINIO_SECURE)
         self.triton_client = httpclient.InferenceServerClient(url=f"{OM_TRITON_IP}:{OM_TRITON_PORT}")
+        self.image_path = request_data.image_path
+        self.result_number = request_data.result_number
+        self.features = self.get_features()
 
     @staticmethod
     def resize_image(img):
@@ -57,8 +61,8 @@ class SimilarMatch:
         mask = np.zeros((1, 1), dtype=np.float32)
         return image, category, mask
 
-    def get_features(self, img_path):
-        image, category, mask = self.preprocess(img_path)
+    def get_features(self):
+        image, category, mask = self.preprocess(self.image_path)
         # 输入集
         inputs = [
             httpclient.InferInput("input__0", image.shape, datatype="FP32"),
@@ -80,7 +84,7 @@ class SimilarMatch:
         return features
 
     @RunTime
-    def match_features(self, features):
+    def match_features(self):
         # 连接milvus
         # 连接milvus
         client = MilvusClient(uri="http://10.1.1.240:19530", db_name="mixi")
@@ -88,8 +92,8 @@ class SimilarMatch:
             search_response = client.search(
                 collection_name="mixi_outfit",  # Replace with the actual name of your collection
                 # Replace with your query vector
-                data=[features[0]],
-                limit=5,  # Max. number of search results to return
+                data=[self.features[0]],
+                limit=self.result_number,  # Max. number of search results to return
                 output_fields=["id", "image_path"],  # Search parameters
             )
             return search_response
@@ -98,8 +102,9 @@ class SimilarMatch:
 
 
 if __name__ == '__main__':
-    service = SimilarMatch()
-    features = service.get_features(img_path="test/2024 SS/MKTS27000.jpg")
-    search_response = service.match_features(features)
+    request_data = SimilarMatchMItem(image_path="test/top/test_top1.jpg", result_number=1)
+
+    service = SimilarMatch(request_data)
+    search_response = service.match_features()
 
     print(json.dumps(search_response, indent=4))