# -*- coding: utf-8 -*-
import sys

import torch
from loguru import logger

sys.path.extend(['..', '../healthCode'])
from utilsK.general import pre_process, post_process, get_return_data


class Model():
    def __init__(self):
        self.par = {'code': {'weights': '../healthCode/weights/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10},
                    'plate': {'weights': '../healthCode/weights/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1},
                    'conf_thres': 0.4,
                    'iou_thres': 0.45,
                    'device': 'cuda:0',
                    'plate_dilate': (0.5, 0.3)
                    }

        ###加载模型
        self.device = torch.device(self.par['device'])
        self.model = torch.jit.load(self.par['code']['weights'])
        self.model_plate = torch.jit.load(self.par['plate']['weights'])


# 防疫模型
class FKModel(Model):
    # def __init__(self):
    #     super().__init__()

    # names, label_arraylist, rainbows, conf_thres, iou_thres
    def process(self, im0, device, img_type):
        try:
            # 预处理
            img, padInfos = pre_process(im0, self.device)
            # 模型推理  code, plate
            if img_type == 'code':
                pred = self.model(img)
            if img_type == 'plate':
                pred = self.model_plate(img)
            boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'],
                                 iou_thres=self.par['iou_thres'],
                                 nc=self.par[img_type]['nc'])  # 后处理
            dataBack = get_return_data(im0, boxes, modelType=img_type, plate_dilate=self.par['plate_dilate'])
            return dataBack
        except Exception as e:
            logger.exception("模型识别异常：{}", e)
            return None
        # for key in dataBack.keys():
        #     if isinstance(dataBack[key], list):
        #         cv2.imwrite('jitimg/%s.jpg' % (key), dataBack[key][0])  ###返回值： dataBack
        '''
        #dataBack= {'type':1,'color':'green','nameImage':'','phoneNumberImage':'','cityImage':'','hsImage':'','plateImage':''}
            
            type:int,	0—行程卡；1—苏康码；2-车牌
            nameImage:	[姓名图像数组,score]
            color: 	green, yellow, red
            cityImage:	[途径地图像数组,score]
            phoneNumberImage:	[手机号图像数组,score]
            IdNumberImage:	[身份证号数组,score]
            hsImage：[核酸检测情况数组,score]	
            plateImage：	[核酸检测情况数组,score]	
        如果没对应目标，返回“空值”
        '''