# -*-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)