|
@@ -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):
|