加入安全绳检测

detectopencvmut0113auta.py

@@ -59,7 +59,7 @@ from utils.general import (LOGGER, Profile, check_file, check_img_size, check_im
                            strtolst, apply_classifier1, apply_classifieruniform, compute_IOU, task, apply_classifierarm)
 from utils.plots import Annotator, colors, save_one_box
 from utils.torch_utils import select_device, smart_inference_mode
-from utils.renwu import newHelmet, newUniform, newFall, Personcount, Arm, Bag, Cross, Extinguisher, Persontre, Bag, Danager
+from utils.renwu import newHelmet, newUniform, newFall, Personcount, Arm, Bag, Cross, Extinguisher, newPersontre, Bag, Danager,CarHelmetBelt
 # from testpool import func1,TestA
 # def my_handler(signum, frame):
 #     exit(0)
@@ -77,20 +77,20 @@ modelnamedir = {'0': 'helmet', '8': 'danager', '10': 'uniform', '14': 'smoke', '
                 '25': 'fall', '29': 'occupancy', '30': 'liquid', '31': 'pressure', '32': 'sleep', '34': 'personcount',
                 '37': 'other', '38': 'duty', '98': 'face', '55': 'oil', '52': 'jingdian', '53': 'rope',
                 '54': 'personcar', '39': 'inspection', '11': 'reflective', '12': 'phone', '66': 'extinguisher',
-                '67': 'tizi', '68': 'menjin', '35': 'arm', '36': 'persontre', '33': 'bag'}
+                '67': 'belt', '68': 'menjin', '35': 'arm', '36': 'persontre', '33': 'bag'}
 modellabeldir = {'0': 'head,person', '8': 'person', '10': 'black_work_clothes,blue_work_clothes,person', '14': 'smoke',
                  '16': 'fire', '21': 'cross', '25': 'fall', '29': 'car', '30': 'liquid', '31': 'pressure',
                  '32': 'sleep', '34': 'personcount', '37': 'other', '38': 'person', '98': 'face', '55': 'oil',
                  '52': 'person,hand,ball', '53': 'rope', '54': 'person', '39': 'person',
                  '11': 'blue,greent,whitet,bluecoat,whitebarcoat,graycoat,baoan,chenyi,other', '12': 'phone',
-                 '66': 'extinguisher', '67': 'person,tizi', '68': 'person', '35': 'barearm,arm',
+                 '66': 'extinguisher', '67': 'person,head,helmet,belt', '68': 'person', '35': 'barearm,arm',
                  '36': 'person,foot,cart,bag,box', '33': 'handbox,handbag'}
 modelalgdir = {}
 personcountdir = {}
 for key, value in modelnamedir.items():
     modelalgdir[value] = key
 taskmap = {'helmet': newHelmet, 'uniform': newUniform, 'fall': newFall, 'personcount': Personcount, 'arm': Arm, 'bag': Bag,
-           'cross': Cross, 'extinguisher': Extinguisher, 'persontre': Persontre, 'bag': Bag, 'danager': Danager}
+           'cross': Cross, 'extinguisher': Extinguisher, 'persontre': newPersontre, 'bag': Bag, 'danager': Danager,'belt':CarHelmetBelt}
 mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
 test = transforms.Compose([transforms.Resize((224, 224)),
                            # transforms.CenterCrop(224),
@@ -775,7 +775,7 @@ if __name__ == '__main__':
         c.execute('UPDATE STREAM set address= (?) where channel =(?)', (address, channel))
     conn.commit()
     cursor = c.execute(
-        "SELECT modelname from CHANGESTREAM where modelname = 'helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname ='duty'  or modelname = 'sleep' or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname= 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall'")
+        "SELECT modelname from CHANGESTREAM where modelname = 'helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname ='duty'  or modelname = 'sleep' or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname= 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall' or modelname = 'belt'")
     # cursor = c.execute("SELECT modelname from CHANGESTREAM where modelname = 'helmet'")
     content = cursor.fetchall()
     cursor = c.execute("SELECT address,channel from STREAM ")
@@ -797,7 +797,7 @@ if __name__ == '__main__':
     print(addcha)
     source = list(source)
     cursor = c.execute(
-        "SELECT modelname from STREAM where (modelname ='helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname = 'duty'  or modelname = 'sleep'  or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname = 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall')")
+        "SELECT modelname from STREAM where (modelname ='helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname = 'duty'  or modelname = 'sleep'  or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname = 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall' or modelname = 'belt')")
     contentm = cursor.fetchall()
     for m in contentm:
         modellist.append(m[0])
@@ -855,7 +855,7 @@ if __name__ == '__main__':
         time.sleep(600)
         task(c, conn, urlt, urla)
         cursor = c.execute(
-            "SELECT modelname from CHANGESTREAM where modelname = 'helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname ='duty'  or modelname = 'sleep' or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname = 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall'")
+            "SELECT modelname from CHANGESTREAM where modelname = 'helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname ='duty'  or modelname = 'sleep' or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname = 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall' or modelname = 'belt'")
         content = cursor.fetchall()
         contentlist = []
         for con in content:
@@ -878,7 +878,7 @@ if __name__ == '__main__':
         source = list(source)
         # source.sort()
         cursor = c.execute(
-            "SELECT modelname from STREAM where (modelname = 'helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname = 'duty'  or modelname = 'sleep' or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname = 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall')")
+            "SELECT modelname from STREAM where (modelname = 'helmet' or modelname = 'smoke' or modelname = 'uniform' or modelname = 'fire' or modelname = 'duty'  or modelname = 'sleep' or modelname='occupancy' or modelname = 'personcar' or modelname = 'phone' or modelname = 'reflective' or modelname = 'extinguisher' or modelname = 'danager' or modelname = 'inspection' or modelname = 'cross' or modelname = 'personcount' or modelname = 'arm' or modelname = 'persontre' or modelname = 'bag' or modelname = 'fall' or modelname = 'belt')")
         contentm = cursor.fetchall()
         for m in contentm:
             modellist.append(m[0])

utils/renwu.py

@@ -942,6 +942,61 @@ class Persontre:
                                 annotator.box_label(cls3_box, None, color=(0, 0, 255))
                                 annotator.box_label(best_cls4_box, None, color=(0, 0, 255))
                                 self.flag = True
+
+class newPersontre:
+
+    def __init__(self):
+
+        self.flag = False
+
+        self.classifier_model = torch.load('/home/h3c/yolo/persontrecls.pt')
+
+    def getflag(self, det, persondet, annotator, fence=0, point=None, names=None, rname=None, num=1, im0=None):
+        self.flag = False
+        cls3 = 1  # 对cart单独一类
+        # cls4 = 4  # 新增加的类别编号（假设为 4）
+        # target_classes = [1, 2, 3]
+        results = Results(annotator.result(), path=None, names=names, boxes=det)
+        boxes = results.boxes
+        print(boxes)
+        person_boxes = []  # 人 车子 脚 以及 指定区域内的脚
+        cls3_boxes = []
+        # cls4_boxes = []
+        # best_cls4_box = []
+        # 处理检测结果
+        for i in range(len(boxes)):
+            cls = int(boxes.cls[i].item())
+            con = boxes.conf[i].item()
+            if cls == 0 and con > 0.1:  # 如果是 "person" 类别  先把对应类别添加上去  先把每一个识别物体输入上去
+                x1, y1, x2, y2 = boxes.xyxy[i].tolist()
+                person_boxes.append([x1, y1, x2, y2])
+            if cls == cls3 and con > 0.1:  # 目标类别 cart
+                x1, y1, x2, y2 = boxes.xyxy[i].tolist()
+                cls3_boxes.append([x1, y1, x2, y2])
+        print(person_boxes)
+        print(cls3_boxes)
+        # 如果检测到 "person" 类别和目标框，计算IoU
+        if person_boxes and cls3_boxes:  # 先查看这俩个类别有没有
+            for i, person_box in enumerate(person_boxes):
+                for j, cls3_box in enumerate(cls3_boxes):
+                    # 判断是否存在交集 而且不能低
+                    iou = calculate_iou(person_box, cls3_box)
+                    print(f"IOU的大小这是第", iou)
+                    if iou > 0.15:  # 用来判断坐 如果iou足够高 就视为是坐
+                        # 加载图像并绘制标注
+                        # 看看这个值大概有多少 好衡量一下
+                        print(f"这个iou大于0.2的 这是第,iou是{iou}")  # 写入介绍
+                        # img = Image.fromarray(cv2.cvtColor(frame,
+                        # cv2.COLOR_BGR2RGB))  # 注意 cap.read 获取到的帧是numpy形式 不能直接配合Draw函数 需要转换
+                        # 创建绘图上下文
+                        # draw = ImageDraw.Draw(img)
+                        # 继续进行绘制操作
+                        # 绘制 "person" 类别框
+                        annotator.box_label(person_box, None, color=(0, 0, 255))
+                        # 绘制目标框（bag, box, cart）
+                        annotator.box_label(cls3_box, None, color=(0, 0, 255))
+                        self.flag = True
+        return self.flag
 class Danager:
     def __init__(self):
         self.flag = False