更新 enums/ModelTypeEnum.py

enums/ModelTypeEnum.py

@@ -16,6 +16,7 @@ from utilsK.illParkingUtils import illParking_postprocess
 from utilsK.pannelpostUtils import pannel_post_process
 from stdc import stdcModel
 from yolov5 import yolov5Model
+from p2pNet import p2NnetModel
 from DMPRUtils.jointUtil import  dmpr_yolo_stdc
 from AI import default_mix
 from ocr import ocrModel
@@ -361,14 +362,14 @@ class ModelType(Enum):
     })
 
     CITY_MANGEMENT_MODEL = ("16", "016", "城管模型", 'cityMangement2', lambda device, gpuName: {
-        'labelnames': [ "车辆", "垃圾", "商贩", "违停","占道经营","裸土","未覆盖裸土" ],
+        'labelnames': [ "车辆", "垃圾", "商贩", "违停","占道经营","裸土","未覆盖裸土","违建" ],
         'postProcess':{
                     'function':dmpr_yolo_stdc,
                     'pars':{
-                            'carCls':0 ,'illCls':6,'scaleRatio':0.5,'border':80,
+                            'carCls':0 ,'illCls':7,'scaleRatio':0.5,'border':80,
-                            #车辆","垃圾","商贩","裸土","占道经营","违停"--->
+                             #"车辆","垃圾","商贩","裸土","占道经营","未覆盖裸土","违建"
-                             #"车辆","垃圾","商贩","违停","占道经营","裸土"
+                             # key:实际训练index     value:展示index
-                             'classReindex':{ 0:0,1:1,2:2,3:6,4:4,5:5,6:3}
+                             'classReindex':{ 0:0,1:1,2:2,7:3,4:4,3:5,5:6,6:7}
                         }
                     },
         'models':[
@@ -376,7 +377,7 @@ class ModelType(Enum):
                 'weight':'../weights/trt/AIlib2/cityMangement3/yolov5_%s_fp16.engine'%(gpuName),
                 'name':'yolov5',
                 'model':yolov5Model,
-                'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3],'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':True, "score_byClass":{"0":0.8,"1":0.4,"2":0.5,"3":0.5  } }
+                'par':{ 'half':True,'device':'cuda:0' ,'conf_thres':0.25,'iou_thres':0.45,'allowedList':[0,1,2,3,4,5,6,7],'segRegionCnt':1, 'trtFlag_det':True,'trtFlag_seg':True, "score_byClass":{"0":0.8,"1":0.4,"2":0.5,"3":0.5  } }
             },
             {
                 'weight':'../weights/pth/AIlib2/cityMangement3/dmpr.pth',
@@ -399,10 +400,10 @@ class ModelType(Enum):
             "name": "post_process",
             "conf_thres": 0.25,
             "iou_thres": 0.45,
-            "classes": 6,
+            "classes": 8,
             "rainbows": COLOR
         },
-        'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,5,6,7,8,9]  ],###控制哪些检测类别显示、输出
+        'detModelpara':[{"id":str(x),"config":{"k1":"v1","k2":"v2"}} for x in [0,1,2,3,4,5,6,7,8,9]  ],###控制哪些检测类别显示、输出
         'segRegionCnt':2,###分割模型结果需要保留的等值线数目
         "pixScale": 1.2,
     })
@@ -664,7 +665,8 @@ class ModelType(Enum):
     })
     TRAFFICFORDSJ_FARM_MODEL = ("27", "027", "交通模型-大数据局", 'highWay2T', lambda device, gpuName: {
         'device': str(device),
-        'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子", "事故", "桥梁外观","设施破损缺失","龙门架","防抛网","标识牌损坏","护栏损坏","钢筋裸露" ],
+        'labelnames': ["行人", "车辆", "纵向裂缝", "横向裂缝", "修补", "网状裂纹", "坑槽", "块状裂纹", "积水", "影子",
+                       "事故", "桥梁外观","设施破损缺失","龙门架","防抛网","标识牌损坏","护栏损坏","钢筋裸露" ],
         'trtFlag_seg': True,
         'trtFlag_det': True,
         'seg_nclass': 3,
@@ -958,6 +960,52 @@ class ModelType(Enum):
 
     })
 
+    CITY_DENSECROWDCOUNT_MODEL = ("30", "304", "密集人群计数", 'DenseCrowdCount', lambda device, gpuName: {
+        'labelnames': ["人群计数"],
+        'device': str(device),
+        'rainbows': COLOR,
+        'models': [
+            {
+                'trtFlag_det': False,
+                'weight': "../weights/pth/AIlib2/DenseCrowd/SHTechA.pth",  ###检测模型路径
+                'vggweight': "../weights/pth/AIlib2/DenseCrowd/vgg16_bn-6c64b313.pth",  ###检测模型路径
+                'name': 'p2pnet',
+                'model': p2NnetModel,
+                'par': {
+                    'device': 'cuda:0',
+                    'row': 2,
+                    'line': 2,
+                    'point_loss_coef': 0.45,
+                    'conf': 0.25,
+                    'gpu_id': 0,
+                    'eos_coef': '0.5',
+                    'set_cost_class': 1,
+                    'set_cost_point': 0.05,
+                    'backbone': 'vgg16_bn'
+                },
+            }],
+    })
+
+    CITY_DENSECROWDESTIMATION_MODEL = ("30", "305", "密集人群密度估计", 'DenseCrowdEstimation', lambda device, gpuName: {
+        'labelnames': ["密度"],
+        'models':
+            [
+                {
+                    'weight': "../weights/pth/AIlib2/DenseCrowd/SHTechA.pth",  ###检测模型路径
+                    'name': 'yolov5',
+                    'model': yolov5Model,
+                    'par': {'half': True, 'device': 'cuda:0', 'conf_thres': 0.50, 'iou_thres': 0.45,
+                            'allowedList': list(range(20)), 'segRegionCnt': 1, 'trtFlag_det': True,
+                            'trtFlag_seg': False, "score_byClass": {"0": 0.50, "1": 0.3, "2": 0.3, "3": 0.3}},
+                }
+
+            ],
+        'postFile': {
+            "rainbows": COLOR
+        },
+
+    })
+
     @staticmethod
     def checkCode(code):
         for model in ModelType: