143 lines
4.7 KiB
Python
143 lines
4.7 KiB
Python
# -*-coding:utf-8-*-
|
|
import os
|
|
import sys
|
|
from logzero import logger
|
|
current_path = os.path.dirname(__file__)
|
|
logger.info(current_path)
|
|
sys.path.append(f"{current_path}/text2image/")
|
|
sys.path.append(f"{current_path}/text2image/BigGAN_utils/")
|
|
import json
|
|
import base64
|
|
from flask import Flask, request, make_response
|
|
import cv2
|
|
# from io import BytesIO
|
|
# import torch
|
|
from mmpose.apis import init_pose_model
|
|
# from text2image.run_text2img import text2image
|
|
# from detection.detection import detector
|
|
# from segmentation.segment_pred import run_seg
|
|
# from ocr.ocr import run_tr
|
|
from backbone.backbone_infer import run_backbone_infer
|
|
|
|
DEVICE = 'cpu'
|
|
|
|
# model_5x = torch.hub.load(f'{current_path}/detection/yolov5/','yolov5x', source='local', pretrained=True)
|
|
# model_5s = torch.hub.load(f'{current_path}/detection/yolov5/','yolov5s', source='local', pretrained=True)
|
|
# 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='cpu') # or device='cuda:0'
|
|
|
|
app=Flask(__name__)
|
|
|
|
# @app.route('/text2image/',methods=["POST"])
|
|
# def run_text2img():
|
|
# if request.method == "POST":
|
|
# text = request.form.get('text')
|
|
# logger.info(f"{text}")
|
|
# img = text2image(text)
|
|
# output_buffer = BytesIO()
|
|
# 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
|
|
|
|
|
|
# @app.route('/detection/', methods=["POST"])
|
|
# def run_detection():
|
|
# 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)
|
|
# 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
|
|
|
|
# @app.route('/ocr/', methods=["POST"])
|
|
# def run_ocr():
|
|
# resp = make_response()
|
|
# 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
|
|
|
|
|
|
# @app.route('/segmentation/', methods=["POST"])
|
|
# def run_segmentation():
|
|
# 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
|
|
|
|
@app.route('/backbone/', methods=["POST"])
|
|
def run_backbone():
|
|
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
|
|
pose, 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
|
|
|
|
|
|
if __name__ == '__main__':
|
|
img = cv2.imread('./1.jpg')
|
|
pose, rst = run_backbone_infer(img, pose_model)
|
|
cv2.imwrite('./1_bb.jpg', rst) |