Idiot-Cultivation-System / src /RemoteDatabase.py
silk-road's picture
Upload 23 files
0117cec verified
# -*- coding: utf-8 -*-
import requests
import pickle
import base64
import json, os
def serialize_feature(feature):
if feature is None:
return None
return base64.encodebytes(pickle.dumps(feature)).decode('ascii')
def deserialize_feature(feature):
if feature is None:
return None
return pickle.loads(base64.decodebytes(feature.encode('ascii')))
class RemoteDatabase:
def __init__(self):
self.url = 'http://110.40.175.218:6007/'
#self.url = 'http://127.0.0.1:6007/'
pass
def top_k_search(self, query_feature, attribute='clip_feature', top_k=15):
url = self.url + 'top_k_search'
query_feature = serialize_feature(query_feature)
response = requests.post(url,
data=json.dumps({"feature":query_feature, "attribute":attribute, "top_k":top_k}),
headers={'Content-Type': 'application/json'})
response.encoding = 'utf-8'
return json.loads(response.text)
def search_by_en_keyword(self, keyword):
url = self.url + 'search_by_en_keyword'
response = requests.post(url,
data=json.dumps({"keyword":keyword}),
headers={'Content-Type': 'application/json'})
response.encoding = 'utf-8'
return json.loads(response.text)
def random_sample(self, n):
url = self.url + 'random_sample'
response = requests.post(url,
data=json.dumps({"number":n}),
headers={'Content-Type': 'application/json'})
response.encoding = 'utf-8'
result = json.loads(response.text)
for r in result:
image_name = r['image_name']
image_name = os.path.basename(image_name)
self.download_file(image_name, os.path.join('temp_images', image_name))
return result
def get_top_founders(self, word="", top_k=20):
url = self.url + 'get_top_founders'
response = requests.post(url,
data=json.dumps({'top_k':top_k, 'word':word}),
headers={'Content-Type': 'application/json'})
response.encoding = 'utf-8'
return json.loads(response.text)
def set_founders(self, word, founder, enforce=False):
url = self.url + 'set_founder'
response = requests.post(url,
data=json.dumps({'founder':founder, 'word':word}),
headers={'Content-Type': 'application/json'})
response.encoding = 'utf-8'
return json.loads(response.text)
def add_data(self, img_data, text_data, img_feature, text_feature):
url = self.url + 'add_data'
img_feature = serialize_feature(img_feature)
text_feature = serialize_feature(text_feature)
response = requests.post(url,
data=json.dumps({"img_data":img_data, "text_data":text_data, "img_feature":img_feature, "text_feature":text_feature}),
headers={'Content-Type': 'application/json'})
response.encoding = 'utf-8'
return json.loads(response.text)
def add_file(self, file_path):
url = self.url + 'add_file'
response = requests.post(url,
data = {},
files = {'file': open(file_path, 'rb')},
stream=True)
response.encoding = 'utf-8'
return json.loads(response.text)
def download_file(self, file_name, save_path):
url = self.url + 'download_file'
params = {'file_name': file_name}
response = requests.post(url, data=params, stream=True)
if response.status_code == 200:
with open(save_path, 'wb') as file:
for chunk in response.iter_content(1024):
if chunk:
file.write(chunk)
if __name__ == '__main__':
db = RemoteDatabase()
db.download_file('xyxmi1lr55pm.jpg', './test.jpg')