zhaojinghao
/
ai_platform_cv
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

修改response结构

This commit is contained in:
zhaojinghao 2022-12-08 14:39:23 +08:00
parent e872462870
commit 7b4e850f9f
1 changed files with 95 additions and 84 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

177
run.py
View File

@ -3,8 +3,9 @@ import os
import sys import sys
from logzero import logger from logzero import logger
# current_path = os.path.dirname(__file__) # for local # current_path = os.path.dirname(__file__) # for local
current_path = "/app" # for docker current_path = "/app" # for docker
logger.info(f"{current_path}") logger.info(f"{current_path}")
sys.path.append(f"{current_path}/text2image/") sys.path.append(f"{current_path}/text2image/")
sys.path.append(f"{current_path}/text2image/BigGAN_utils/") sys.path.append(f"{current_path}/text2image/BigGAN_utils/")
@ -24,20 +25,27 @@ from ocr.ocr import run_tr
from backbone.backbone_infer import run_backbone_infer from backbone.backbone_infer import run_backbone_infer
from mnist.mnist_torch import run_mnist_infer, CNN from mnist.mnist_torch import run_mnist_infer, CNN
DEVICE = 'cpu'
model_5x = torch.hub.load(f'{current_path}/detection/yolov5/','custom', path=f'{current_path}/detection/models/yolov5x.pt', source='local') DEVICE = 'cpu'
model_5s = torch.hub.load(f'{current_path}/detection/yolov5/','custom', path=f'{current_path}/detection/models/yolov5s.pt', source='local') TEXT = "text"
BASE64_IMG = "base64_img"
model_5x = torch.hub.load(f'{current_path}/detection/yolov5/', 'custom',
path=f'{current_path}/detection/models/yolov5x.pt', source='local')
model_5s = torch.hub.load(f'{current_path}/detection/yolov5/', 'custom',
path=f'{current_path}/detection/models/yolov5s.pt', source='local')
model_seg = torch.load(f'{current_path}/segmentation/models/best_model_pvgc.pth', map_location=DEVICE) model_seg = torch.load(f'{current_path}/segmentation/models/best_model_pvgc.pth', map_location=DEVICE)
pose_config_file = f'{current_path}/backbone/associative_embedding_hrnet_w32_coco_512x512.py' pose_config_file = f'{current_path}/backbone/associative_embedding_hrnet_w32_coco_512x512.py'
pose_ckpt_file = f'{current_path}/backbone/models/hrnet_w32_coco_512x512-bcb8c247_20200816.pth' pose_ckpt_file = f'{current_path}/backbone/models/hrnet_w32_coco_512x512-bcb8c247_20200816.pth'
pose_model = init_pose_model(pose_config_file, pose_ckpt_file, device=DEVICE) pose_model = init_pose_model(pose_config_file, pose_ckpt_file, device=DEVICE)
mnist_model = torch.load(f'{current_path}/mnist/models/MNIST_torch.pth', map_location=DEVICE) mnist_model = torch.load(f'{current_path}/mnist/models/MNIST_torch.pth', map_location=DEVICE)
app=Flask(__name__) app = Flask(__name__)
@app.route('/text2image/',methods=["POST"])
@app.route('/text2image/', methods=["POST"])
def run_text2img(): def run_text2img():
resp_info = dict()
if request.method == "POST": if request.method == "POST":
text = request.form.get('text') text = request.form.get('text')
logger.info(f"{text}") logger.info(f"{text}")
@ -46,127 +54,130 @@ def run_text2img():
img.save(output_buffer, format='png') img.save(output_buffer, format='png')
byte_data = output_buffer.getvalue() byte_data = output_buffer.getvalue()
b64_code = base64.b64encode(byte_data).decode('utf-8') b64_code = base64.b64encode(byte_data).decode('utf-8')
resp = make_response(b64_code) resp_info["code"] = 200
resp.status_code = 200 resp_info["data"] = b64_code
return resp resp_info["dtype"] = BASE64_IMG
else: resp = make_response(resp_info)
resp = make_response() resp.status_code = 200
resp.status_code=405 return resp
return resp
@app.route('/detection/', methods=["POST"]) @app.route('/detection/', methods=["POST"])
def run_detection(): def run_detection():
resp_info = dict()
if request.method == "POST": if request.method == "POST":
img = request.files.get('image') img = request.files.get('image')
model_type = request.form.get('model_type') model_type = request.form.get('model_type')
try: try:
img = cv2.imread(img) img = cv2.imread(img)
except: except Exception as e:
resp = make_response() resp_info["msg"] = e
resp.status_code = 406 resp_info["code"] = 406
return resp
if model_type.lower().strip() == 'yolov5x':
rst, _ = detector(img, model_5x)
else: else:
rst, _ = detector(img, model_5s) if model_type.lower().strip() == 'yolov5x':
logger.info(rst.shape) rst, _ = detector(img, model_5x)
img_str = cv2.imencode('.png', rst)[1].tobytes() else:
b64_code = base64.b64encode(img_str).decode('utf-8') rst, _ = detector(img, model_5s)
resp = make_response(b64_code) logger.info(rst.shape)
resp.status_code = 200 img_str = cv2.imencode('.png', rst)[1].tobytes()
return b64_code b64_code = base64.b64encode(img_str).decode('utf-8')
else: resp_info["code"] = 200
resp = make_response() resp_info["data"] = b64_code
resp.status_code=405 resp_info["dtype"] = BASE64_IMG
return resp resp = make_response(resp_info)
resp.status_code = 200
return resp
@app.route('/ocr/', methods=["POST"]) @app.route('/ocr/', methods=["POST"])
def run_ocr(): def run_ocr():
resp = make_response() resp_info = dict()
if request.method == "POST": if request.method == "POST":
img = request.files.get('image') img = request.files.get('image')
try: try:
img = cv2.imread(img) img = cv2.imread(img)
except: except Exception as e:
resp.status_code = 406 resp_info["msg"] = e
return resp resp_info["code"] = 406
text = run_tr(img) else:
resp.status_code = 200 text = run_tr(img)
resp.data = json.dumps({'result':text}) resp_info["code"] = 200
return resp resp_info["data"] = text
else: resp_info["dtype"] = TEXT
resp.status_code=405 resp = make_response(resp_info)
return resp resp.status_code = 200
return resp
@app.route('/segmentation/', methods=["POST"]) @app.route('/segmentation/', methods=["POST"])
def run_segmentation(): def run_segmentation():
resp_info = dict()
if request.method == "POST": if request.method == "POST":
img_upload = request.files.get('image') img_upload = request.files.get('image')
try: try:
img = cv2.imread(img_upload) img = cv2.imread(img_upload)
except: except Exception as e:
resp = make_response() resp_info["msg"] = e
resp.status_code = 406 resp_info["code"] = 406
return resp else:
result = run_seg(img, model_seg) result = run_seg(img, model_seg)
img_str = cv2.imencode('.png', result)[1].tobytes() img_str = cv2.imencode('.png', result)[1].tobytes()
b64_code = base64.b64encode(img_str).decode('utf-8') b64_code = base64.b64encode(img_str).decode('utf-8')
resp = make_response(b64_code) resp_info["code"] = 200
resp.status_code = 200 resp_info["data"] = b64_code
return resp resp_info["dtype"] = BASE64_IMG
else: resp = make_response(resp_info)
resp = make_response() resp.status_code = 200
resp.status_code=405 return resp
return resp
@app.route('/backbone/', methods=["POST"]) @app.route('/backbone/', methods=["POST"])
def run_backbone(): def run_backbone():
resp_info = dict()
if request.method == "POST": if request.method == "POST":
img_upload = request.files.get('image') img_upload = request.files.get('image')
try: try:
img = cv2.imread(img_upload) img = cv2.imread(img_upload)
except: except Exception as e:
resp = make_response() resp_info["msg"] = e
resp.status_code = 406 resp_info["code"] = 406
return resp else:
_, result = run_backbone_infer(img, pose_model) _, result = run_backbone_infer(img, pose_model)
img_str = cv2.imencode('.png', result)[1].tobytes() img_str = cv2.imencode('.png', result)[1].tobytes()
b64_code = base64.b64encode(img_str).decode('utf-8') b64_code = base64.b64encode(img_str).decode('utf-8')
resp = make_response(b64_code) resp_info["code"] = 200
resp.status_code = 200 resp_info["data"] = b64_code
return resp resp_info["dtype"] = BASE64_IMG
else: resp = make_response(resp_info)
resp = make_response() resp.status_code = 200
resp.status_code=405 return resp
return resp
@app.route('/mnist/', methods=["POST"]) @app.route('/mnist/', methods=["POST"])
def run_mnist(): def run_mnist():
resp_info = dict()
if request.method == "POST": if request.method == "POST":
img_upload = request.files.get('image') img_upload = request.files.get('image')
try: try:
img = cv2.imread(img_upload, 1) img = cv2.imread(img_upload, 1)
# 使用全局阈值，降噪 # 使用全局阈值，降噪
ret,th1 = cv2.threshold(img,127,255,cv2.THRESH_BINARY) ret, th1 = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
# 把opencv图像转化为PIL图像 # 把opencv图像转化为PIL图像
im = Image.fromarray(cv2.cvtColor(th1,cv2.COLOR_BGR2RGB)) im = Image.fromarray(cv2.cvtColor(th1, cv2.COLOR_BGR2RGB))
# 灰度化 # 灰度化
im = im.convert('L') im = im.convert('L')
Im = im.resize((28, 28), Image.ANTIALIAS) Im = im.resize((28, 28), Image.ANTIALIAS)
except: except Exception as e:
resp = make_response() resp_info["msg"] = e
resp.status_code = 406 resp_info["code"] = 406
return resp else:
result = run_mnist_infer(Im, mnist_model) result = run_mnist_infer(Im, mnist_model)
resp = make_response(str(result)) resp_info["code"] = 200
resp.status_code = 200 resp_info["data"] = str(result)
return resp resp_info["dtype"] = TEXT
else: resp = make_response(resp_info)
resp = make_response() resp.status_code = 200
resp.status_code=405 return resp
return resp
if __name__ == '__main__': if __name__ == '__main__':