空港防疫算法交互
Du kan inte välja fler än 25 ämnen Ämnen måste starta med en bokstav eller siffra, kan innehålla bindestreck ('-') och vara max 35 tecken långa.

66 lines
2.6KB

  1. # -*- coding: utf-8 -*-
  2. import sys
  3. import torch
  4. from loguru import logger
  5. sys.path.extend(['..', '../healthCode'])
  6. from utilsK.general import pre_process, post_process, get_return_data
  7. class Model():
  8. def __init__(self):
  9. self.par = {'code': {'weights': '../healthCode/weights/health_yolov5s_v3.jit', 'img_type': 'code', 'nc': 10},
  10. 'plate': {'weights': '../healthCode/weights/plate_yolov5s_v3.jit', 'img_type': 'plate', 'nc': 1},
  11. 'conf_thres': 0.4,
  12. 'iou_thres': 0.45,
  13. 'device': 'cuda:0',
  14. 'plate_dilate': (0.5, 0.3)
  15. }
  16. ###加载模型
  17. self.device = torch.device(self.par['device'])
  18. self.model = torch.jit.load(self.par['code']['weights'])
  19. self.model_plate = torch.jit.load(self.par['plate']['weights'])
  20. # 防疫模型
  21. class FKModel(Model):
  22. # def __init__(self):
  23. # super().__init__()
  24. # names, label_arraylist, rainbows, conf_thres, iou_thres
  25. def process(self, im0, device, img_type):
  26. try:
  27. # 预处理
  28. img, padInfos = pre_process(im0, self.device)
  29. # 模型推理 code, plate
  30. if img_type == 'code':
  31. pred = self.model(img)
  32. if img_type == 'plate':
  33. pred = self.model_plate(img)
  34. boxes = post_process(pred, padInfos, self.device, conf_thres=self.par['conf_thres'],
  35. iou_thres=self.par['iou_thres'],
  36. nc=self.par[img_type]['nc']) # 后处理
  37. dataBack = get_return_data(im0, boxes, modelType=img_type, plate_dilate=self.par['plate_dilate'])
  38. return dataBack
  39. except Exception as e:
  40. logger.exception("模型识别异常:{}", e)
  41. return None
  42. # for key in dataBack.keys():
  43. # if isinstance(dataBack[key], list):
  44. # cv2.imwrite('jitimg/%s.jpg' % (key), dataBack[key][0]) ###返回值: dataBack
  45. '''
  46. #dataBack= {'type':1,'color':'green','nameImage':'','phoneNumberImage':'','cityImage':'','hsImage':'','plateImage':''}
  47. type:int, 0—行程卡;1—苏康码;2-车牌
  48. nameImage: [姓名图像数组,score]
  49. color: green, yellow, red
  50. cityImage: [途径地图像数组,score]
  51. phoneNumberImage: [手机号图像数组,score]
  52. IdNumberImage: [身份证号数组,score]
  53. hsImage:[核酸检测情况数组,score]
  54. plateImage: [核酸检测情况数组,score]
  55. 如果没对应目标,返回“空值”
  56. '''