# -*- coding: utf-8 -*-
from flask import request, Flask
import pdb
import numpy as np
import json
import time
import cv2, base64
import argparse
import sys, os
import torch
from PIL import Image
import datetime
from torchvision import transforms
from models.experimental import attempt_load
from utils.torch_utils import select_device
# from models.modeling import VisionTransformer, CONFIGS
from gevent.pywsgi import WSGIServer
sys.path.insert(0, ".")

import detect
import logging.config
from skywalking import agent, config
SW_SERVER = os.environ.get('SW_AGENT_COLLECTOR_BACKEND_SERVICES')
SW_SERVICE_NAME = os.environ.get('SW_AGENT_NAME')
if SW_SERVER and SW_SERVICE_NAME:
    config.init() #采集服务的地址，给自己的服务起个名称
    #config.init(collector="123.60.56.51:11800", service='ieemoo-ai-search') #采集服务的地址，给自己的服务起个名称
    agent.start()
def setup_logging(path):
    if os.path.exists(path):
        with open(path, 'r') as f:
            config = json.load(f)
            logging.config.dictConfig(config)
            logger = logging.getLogger("root")
    return logger
logger = setup_logging('utils/logging.json')

app = Flask(__name__)
app.use_reloader=False

parser = argparse.ArgumentParser()
#parser.add_argument('--weights', nargs='+', type=str, default='runs/zhanting/yolov5m_finetune/exp4/weights/best.pt', help='model.pt path(s)')
parser.add_argument('--weights', nargs='+', type=str, default='../module/ieemoo-ai-zhanting/model/now/best.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='../module/ieemoo-ai-zhanting/imgs/1.jpg', help='source')  # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.60, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', type=bool, default=True, help='display results')
parser.add_argument('--save-txt', type=bool, default=True, help='save results to *.txt')
parser.add_argument('--save-conf', type=bool, default=True, help='save confidences in --save-txt labels')
parser.add_argument('--nosave', type=bool, default=True, help='do not save images/videos')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', type=bool, default=True, help='class-agnostic NMS')
parser.add_argument('--augment', type=bool, default=True, help='augmented inference')
parser.add_argument('--update', type=bool, default=True, help='update all models')
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
parser.add_argument('--name', default='exp', help='save results to project/name')
parser.add_argument('--exist-ok', type=bool, default=True, help='existing project/name ok, do not increment')

#opt = parser.parse_args()
opt, unknown = parser.parse_known_args()
def init_model(opt):
     device = select_device(opt.device)
     #half = device.type != 'cpu'
     model = attempt_load(opt.weights, map_location=device)
     stride = int(model.stride.max())
     #if half:
     model.half()
     model.eval()
     return model, stride
model,stride = init_model(opt)

@app.route("/zhanting", methods=['POST'])
def get_isempty():
    data = request.get_data()
    ip = request.remote_addr
    print('------ ip = %s ------' % ip)
    json_data = json.loads(data.decode("utf-8"))
    pic = json_data.get("pic")
    result = {"success": "false",
              "rst_cls": '-1',
              }
    try:
        imgdata = base64.b64decode(pic)
        image_path = '../module/ieemoo-ai-zhanting/imgs/1.jpg'
        file = open(image_path, 'wb')
        file.write(imgdata)
        img = cv2.imread(image_path)
        #pdb.set_trace()
        if img.shape[0] == 1280:
            site = np.array([[[0, 1280],[0, 671], [300, 390], [740, 390], [1024, 635], [1024,1280]]], dtype=np.int32)
        else:
            site = np.array([[[0, 1024],[0, 571], [313, 365], [949, 367], [1277, 596], [1280, 1024]]], dtype=np.int32)
        im = np.zeros(img.shape[:2], dtype="uint8")
        cv2.polylines(im, site, 1, 255)
        cv2.fillPoly(im, site, 255)
        mask = im
        masked = cv2.bitwise_or(img, img, mask=mask)
        img0 = masked
        cv2.imwrite("../module/ieemoo-ai-zhanting/imgs/1.jpg",img0)
        pred = detect.detect(opt, model, stride)
        logger.info(pred)
    except Exception as e:
        logger.warning(e)
        return pred
    return pred

if __name__ == "__main__":
    app.run(host='192.168.1.142', port=8000)