|
- # -*- 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]
- 如果没对应目标，返回“空值”
- '''
|
