修改response结构
This commit is contained in:
zhaojinghao
parent
e872462870
commit
7b4e850f9f
179
run.py
179
run.py
|
|
@ -3,8 +3,9 @@ import os
|
|||
import sys
|
||||
from logzero import logger
|
||||
|
||||
|
||||
# current_path = os.path.dirname(__file__) # for local
|
||||
current_path = "/app" # for docker
|
||||
current_path = "/app" # for docker
|
||||
logger.info(f"{current_path}")
|
||||
sys.path.append(f"{current_path}/text2image/")
|
||||
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 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')
|
||||
model_5s = torch.hub.load(f'{current_path}/detection/yolov5/','custom', path=f'{current_path}/detection/models/yolov5s.pt', source='local')
|
||||
DEVICE = 'cpu'
|
||||
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)
|
||||
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_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)
|
||||
|
||||
app=Flask(__name__)
|
||||
app = Flask(__name__)
|
||||
|
||||
@app.route('/text2image/',methods=["POST"])
|
||||
|
||||
@app.route('/text2image/', methods=["POST"])
|
||||
def run_text2img():
|
||||
resp_info = dict()
|
||||
if request.method == "POST":
|
||||
text = request.form.get('text')
|
||||
logger.info(f"{text}")
|
||||
|
|
@ -46,128 +54,131 @@ def run_text2img():
|
|||
img.save(output_buffer, format='png')
|
||||
byte_data = output_buffer.getvalue()
|
||||
b64_code = base64.b64encode(byte_data).decode('utf-8')
|
||||
resp = make_response(b64_code)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
else:
|
||||
resp = make_response()
|
||||
resp.status_code=405
|
||||
return resp
|
||||
resp_info["code"] = 200
|
||||
resp_info["data"] = b64_code
|
||||
resp_info["dtype"] = BASE64_IMG
|
||||
resp = make_response(resp_info)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
|
||||
@app.route('/detection/', methods=["POST"])
|
||||
def run_detection():
|
||||
resp_info = dict()
|
||||
if request.method == "POST":
|
||||
img = request.files.get('image')
|
||||
model_type = request.form.get('model_type')
|
||||
try:
|
||||
img = cv2.imread(img)
|
||||
except:
|
||||
resp = make_response()
|
||||
resp.status_code = 406
|
||||
return resp
|
||||
if model_type.lower().strip() == 'yolov5x':
|
||||
rst, _ = detector(img, model_5x)
|
||||
except Exception as e:
|
||||
resp_info["msg"] = e
|
||||
resp_info["code"] = 406
|
||||
else:
|
||||
rst, _ = detector(img, model_5s)
|
||||
logger.info(rst.shape)
|
||||
img_str = cv2.imencode('.png', rst)[1].tobytes()
|
||||
b64_code = base64.b64encode(img_str).decode('utf-8')
|
||||
resp = make_response(b64_code)
|
||||
resp.status_code = 200
|
||||
return b64_code
|
||||
else:
|
||||
resp = make_response()
|
||||
resp.status_code=405
|
||||
return resp
|
||||
if model_type.lower().strip() == 'yolov5x':
|
||||
rst, _ = detector(img, model_5x)
|
||||
else:
|
||||
rst, _ = detector(img, model_5s)
|
||||
logger.info(rst.shape)
|
||||
img_str = cv2.imencode('.png', rst)[1].tobytes()
|
||||
b64_code = base64.b64encode(img_str).decode('utf-8')
|
||||
resp_info["code"] = 200
|
||||
resp_info["data"] = b64_code
|
||||
resp_info["dtype"] = BASE64_IMG
|
||||
resp = make_response(resp_info)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
|
||||
@app.route('/ocr/', methods=["POST"])
|
||||
def run_ocr():
|
||||
resp = make_response()
|
||||
resp_info = dict()
|
||||
if request.method == "POST":
|
||||
img = request.files.get('image')
|
||||
try:
|
||||
img = cv2.imread(img)
|
||||
except:
|
||||
resp.status_code = 406
|
||||
return resp
|
||||
text = run_tr(img)
|
||||
resp.status_code = 200
|
||||
resp.data = json.dumps({'result':text})
|
||||
return resp
|
||||
else:
|
||||
resp.status_code=405
|
||||
return resp
|
||||
except Exception as e:
|
||||
resp_info["msg"] = e
|
||||
resp_info["code"] = 406
|
||||
else:
|
||||
text = run_tr(img)
|
||||
resp_info["code"] = 200
|
||||
resp_info["data"] = text
|
||||
resp_info["dtype"] = TEXT
|
||||
resp = make_response(resp_info)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
|
||||
@app.route('/segmentation/', methods=["POST"])
|
||||
def run_segmentation():
|
||||
resp_info = dict()
|
||||
if request.method == "POST":
|
||||
img_upload = request.files.get('image')
|
||||
try:
|
||||
img = cv2.imread(img_upload)
|
||||
except:
|
||||
resp = make_response()
|
||||
resp.status_code = 406
|
||||
return resp
|
||||
result = run_seg(img, model_seg)
|
||||
img_str = cv2.imencode('.png', result)[1].tobytes()
|
||||
b64_code = base64.b64encode(img_str).decode('utf-8')
|
||||
resp = make_response(b64_code)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
else:
|
||||
resp = make_response()
|
||||
resp.status_code=405
|
||||
return resp
|
||||
except Exception as e:
|
||||
resp_info["msg"] = e
|
||||
resp_info["code"] = 406
|
||||
else:
|
||||
result = run_seg(img, model_seg)
|
||||
img_str = cv2.imencode('.png', result)[1].tobytes()
|
||||
b64_code = base64.b64encode(img_str).decode('utf-8')
|
||||
resp_info["code"] = 200
|
||||
resp_info["data"] = b64_code
|
||||
resp_info["dtype"] = BASE64_IMG
|
||||
resp = make_response(resp_info)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
|
||||
@app.route('/backbone/', methods=["POST"])
|
||||
def run_backbone():
|
||||
resp_info = dict()
|
||||
if request.method == "POST":
|
||||
img_upload = request.files.get('image')
|
||||
try:
|
||||
img = cv2.imread(img_upload)
|
||||
except:
|
||||
resp = make_response()
|
||||
resp.status_code = 406
|
||||
return resp
|
||||
_, result = run_backbone_infer(img, pose_model)
|
||||
img_str = cv2.imencode('.png', result)[1].tobytes()
|
||||
b64_code = base64.b64encode(img_str).decode('utf-8')
|
||||
resp = make_response(b64_code)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
else:
|
||||
resp = make_response()
|
||||
resp.status_code=405
|
||||
return resp
|
||||
except Exception as e:
|
||||
resp_info["msg"] = e
|
||||
resp_info["code"] = 406
|
||||
else:
|
||||
_, result = run_backbone_infer(img, pose_model)
|
||||
img_str = cv2.imencode('.png', result)[1].tobytes()
|
||||
b64_code = base64.b64encode(img_str).decode('utf-8')
|
||||
resp_info["code"] = 200
|
||||
resp_info["data"] = b64_code
|
||||
resp_info["dtype"] = BASE64_IMG
|
||||
resp = make_response(resp_info)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
|
||||
@app.route('/mnist/', methods=["POST"])
|
||||
def run_mnist():
|
||||
resp_info = dict()
|
||||
if request.method == "POST":
|
||||
img_upload = request.files.get('image')
|
||||
try:
|
||||
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图像
|
||||
im = Image.fromarray(cv2.cvtColor(th1,cv2.COLOR_BGR2RGB))
|
||||
im = Image.fromarray(cv2.cvtColor(th1, cv2.COLOR_BGR2RGB))
|
||||
# 灰度化
|
||||
im = im.convert('L')
|
||||
Im = im.resize((28, 28), Image.ANTIALIAS)
|
||||
except:
|
||||
resp = make_response()
|
||||
resp.status_code = 406
|
||||
return resp
|
||||
result = run_mnist_infer(Im, mnist_model)
|
||||
resp = make_response(str(result))
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
else:
|
||||
resp = make_response()
|
||||
resp.status_code=405
|
||||
return resp
|
||||
except Exception as e:
|
||||
resp_info["msg"] = e
|
||||
resp_info["code"] = 406
|
||||
else:
|
||||
result = run_mnist_infer(Im, mnist_model)
|
||||
resp_info["code"] = 200
|
||||
resp_info["data"] = str(result)
|
||||
resp_info["dtype"] = TEXT
|
||||
resp = make_response(resp_info)
|
||||
resp.status_code = 200
|
||||
return resp
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
app.run(host='0.0.0.0', port='8902')
|
||||
app.run(host='0.0.0.0', port='8902')
|
||||
|
|
|
|||
Loading…
Reference in New Issue