new backend

This commit is contained in:
zjt 2023-10-29 13:54:25 +08:00
parent 040aba1654
commit bbfe75ff14
2 changed files with 23 additions and 12 deletions

View File

@ -1,10 +1,11 @@
import base64 import base64
from io import BytesIO
import numpy as np import numpy as np
import cv2 import cv2
from segment_anything import sam_model_registry, SamPredictor from segment_anything import sam_model_registry, SamPredictor
import torch import torch
import time import time
from flask import Flask, request from flask import Flask, Response, request
# check torch version and if cuda is available # check torch version and if cuda is available
print(torch.__version__) print(torch.__version__)
@ -24,7 +25,9 @@ app = Flask(__name__)
def base64_to_image(base64_code): def base64_to_image(base64_code):
img_data = base64.b64decode(base64_code) img_data = base64.b64decode(base64_code)
img_array = np.fromstring(img_data, np.uint8) img_array = np.fromstring(img_data, np.uint8)
img = cv2.imdecode(img_array, cv2.COLOR_RGB2BGR) print(img_array)
img = cv2.imdecode(img_array, -1)
print(img)
return img return img
@ -38,16 +41,19 @@ def index():
if (content_type == 'application/json'): if (content_type == 'application/json'):
print('start calculate embedding') print('start calculate embedding')
base64data = request.get_json()['imgurl'] base64data = request.get_json()['imgurl']
print(base64data)
image = base64_to_image(base64data) image = base64_to_image(base64data)
predictor.set_image(image) predictor.set_image(image)
image_embedding = predictor.get_image_embedding().cpu().numpy() image_embedding = predictor.get_image_embedding().cpu().numpy()
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
print(time.time()) byte_io = BytesIO()
return image_embedding np.save(byte_io, image_embedding)
byte_io.seek(0)
np.save("array.npy", image_embedding)
response = Response(byte_io, mimetype="application/octet-stream")
response.headers["Content-Length"] = 4194432
return response
else: else:
return 'Content-Type not supported!' return 'Content-Type not supported!'
# app.run(port=2020, host="0.0.0.0", debug=True)
print('server starts working') print('server starts working')

File diff suppressed because one or more lines are too long