61 lines
2.1 KiB
Python
61 lines
2.1 KiB
Python
import json
|
||
import os
|
||
import urllib3
|
||
import urllib.parse
|
||
|
||
|
||
class IndexHttpClient():
|
||
"""索引库客户端，使用 urllib3 连接，使用 urllib.parse 进行 url 编码"""
|
||
def __init__(self, host: str, port: int):
|
||
self.__host = host
|
||
self.__port = port
|
||
self.__http = urllib3.PoolManager()
|
||
self.__headers = {"Content-type": "application/json"}
|
||
|
||
def url(self):
|
||
return "http://{}:{}".format(self.__host, self.__port)
|
||
|
||
def new_index(self,
|
||
image_list_path: str,
|
||
index_root_path: str,
|
||
index_method="HNSW32",
|
||
force=False):
|
||
"""新建 重建 库"""
|
||
if index_method not in ["HNSW32", "FLAT", "IVF"]:
|
||
raise Exception(
|
||
"index_method 必须是 HNSW32, FLAT, IVF，实际值为：{}".format(
|
||
index_method))
|
||
params = {"image_list_path":image_list_path, \
|
||
"index_root_path":index_root_path, \
|
||
"index_method":index_method, \
|
||
"force":force}
|
||
return self.__post(self.url() + "/new_index?", params)
|
||
|
||
def open_index(self, index_root_path: str, image_list_path: str):
|
||
"""打开库"""
|
||
params = {
|
||
"index_root_path": index_root_path,
|
||
"image_list_path": image_list_path
|
||
}
|
||
return self.__post(self.url() + "/open_index?", params)
|
||
|
||
def update_index(self, image_list_path: str, index_root_path: str):
|
||
"""更新索引库"""
|
||
params = {"image_list_path":image_list_path, \
|
||
"index_root_path":index_root_path}
|
||
return self.__post(self.url() + "/update_index?", params)
|
||
|
||
def __post(self, url: str, params: dict):
|
||
"""发送 url 并接收数据"""
|
||
http = self.__http
|
||
encode_params = urllib.parse.urlencode(params)
|
||
get_url = url + encode_params
|
||
req = http.request("GET", get_url, headers=self.__headers)
|
||
result = json.loads(req.data)
|
||
if isinstance(result, str):
|
||
result = eval(result)
|
||
msg = result["error_message"]
|
||
if msg != None and len(msg) == 0:
|
||
msg = None
|
||
return msg
|