Browse Source

fix v1.0.2

sun-pc 2 months ago
parent
commit
41394fda3f
2 changed files with 9 additions and 12 deletions
  1. 2 2
      build.sh
  2. 7 10
      src/main.py

+ 2 - 2
build.sh

@@ -1,8 +1,8 @@
 #!/bin/bash
 
 # 命名镜像
-local_imge="search-image:v1.0.0"
-repository_image="registry.cn-chengdu.aliyuncs.com/infish/search-image:v1.0.0"
+local_imge="search-image:v1.0.2"
+repository_image="registry.cn-chengdu.aliyuncs.com/infish/search-image:v1.0.2"
 
 # 删除本地已存在的镜像
 docker rmi $repository_image

+ 7 - 10
src/main.py

@@ -109,15 +109,12 @@ async def upload_images(
 async def search_images(
     table_name: str = None,
     partition_names: list = [],
-    # image: UploadFile = File(...),
     url: str = None,
-    limit: int = TOP_K
+    limit: int = TOP_K,
+    min_score: float = None,
+    max_score: float = None
 ):
     try:
-        # content = await image.read()
-        # img_path = os.path.join(UPLOAD_PATH, image.filename)
-        # with open(img_path, "wb+") as f:
-        #     f.write(content)
         if url is not None:
             img_path = os.path.join(UPLOAD_PATH, os.path.basename(url))
             urlretrieve(url, img_path)
@@ -125,9 +122,9 @@ async def search_images(
         list = []
         for hits in res:
             for hit in hits:
-                list.append({"im_hash":hit.id,"product_id":hit.entity.product_id,"score":hit.distance})
-        # res = dict(zip(paths, distances))
-        # res = sorted(res.items(), key=lambda item: item[1])
+                score = hit.distance
+                if (min_score is None or score >= min_score) and (max_score is None or score <= max_score):
+                    list.append({"im_hash":hit.id,"product_id":hit.entity.product_id,"score":score})
         LOGGER.info("Successfully searched similar images!")
         return list
     except Exception as e:
@@ -151,7 +148,7 @@ async def delete_record(table_name: str = None,partition_name: str=None,im_hash:
     except Exception as e:
         LOGGER.error(e)
         return {'status': False, 'msg': e}
-    
+
 
 @app.post('/collection/count')
 async def count_images(table_name: str = None):