Эх сурвалжийг харах

1. 增加时间段识别 2. 删除一个识别围栏

leon 2 өдөр өмнө
parent
commit
7a1f0e0197
4 өөрчлөгдсөн 107 нэмэгдсэн , 57 устгасан
  1. 4 4
      infer.py
  2. 75 48
      main.py
  3. 3 0
      stream.py
  4. 25 5
      test.py

+ 4 - 4
infer.py

@@ -71,7 +71,7 @@ class ImageTransformer:
 class AreaManage(object):
     def __init__(self):
         self.door_areas = [
-            [(934, 143), (958, 130), (961, 165), (936, 181)],
+            # [(934, 143), (958, 130), (961, 165), (936, 181)],
             [(564, 399), (682, 432), (684, 528), (574, 493)]]
         
         self.door_perspective_matrix = []
@@ -84,7 +84,7 @@ class AreaManage(object):
 
         # 预设值的三个门对应的人的活动区域
         self.person_areas = [
-            [(900, 152),  (914, 84), (637, 20), (604, 73)],
+            # [(900, 152),  (914, 84), (637, 20), (604, 73)],
             [(860, 200), (958, 226), (687, 432), (586, 395)]]
         
         self.person_areas_polygon = [Polygon(area) for area in self.person_areas]
@@ -102,7 +102,8 @@ class AreaManage(object):
         
     def inner_area(self, person_polygon):
         for i in range(len(self.person_areas_polygon)):
-            if person_polygon.intersection(self.person_areas_polygon[i]).area / person_polygon.area > 0.5:
+            overlap_ratop = person_polygon.intersection(self.person_areas_polygon[i]).area / person_polygon.area
+            if overlap_ratop > 0.4:
                 return i
         return -1
 
@@ -196,7 +197,6 @@ class DoorInference(object):
     def __call__(self, image):
         inner_person = {0:[], 1:[], 2:[]}
         person_boxes, person_confs = self.person_detect(image)
-        
         for person_box, person_conf in zip(person_boxes, person_confs):
             person_polygon = Polygon(person_box)
             idx = self.am.inner_area(person_polygon)

+ 75 - 48
main.py

@@ -5,6 +5,22 @@ import cv2
 from stream import StreamCapture
 from infer import DoorInference
 from logger import logger
+import time
+from datetime import datetime
+
+def is_time_in_period(start_time_str, end_time_str):
+    # 解析开始和结束时间字符串为time对象
+    start_time = datetime.strptime(start_time_str, "%H:%M").time()
+    end_time = datetime.strptime(end_time_str, "%H:%M").time()
+    current_time = datetime.now().time()
+    
+    # 判断时间段是否跨天
+    if start_time <= end_time:
+        # 时间段在同一天
+        return start_time <= current_time <= end_time
+    else:
+        # 时间段跨天(如22:00到次日02:00)
+        return current_time >= start_time or current_time <= end_time
 
 async def upload_image(session, url, payload, files):
     try:
@@ -35,57 +51,68 @@ async def process_stream():
     channel = '45'
     stream = StreamCapture(ip, channel)
     
+    start_det_time = "06:00"
+    end_det_time = "20:00"
+
     posttime = time.time() - 30
     async with aiohttp.ClientSession() as session:
-        for frame, ret in stream():
-            if not ret: 
-                continue
-            if time.time() - posttime < 30:
+        while True:
+            if not is_time_in_period(start_det_time, end_det_time):
+                time.sleep(10)
                 continue
-            image = frame.copy()
-            result = instance(image)
-            if len(result) > 0:
-                try:
-                    posttime = time.time()
-                    videoTime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
-                    fileTime = time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())
-                    filename = fileTime + ".jpg"
-                    filenameori = fileTime + "det.jpg"
-                    logger.info(videoTime)
-                    logger.info(result)
-                    
-                    for res in result:
-                        cv2.rectangle(image, tuple(map(int, (res.left, res.top))), 
-                                      tuple(map(int, (res.right, res.bottom))), (255, 0, 0), 4)
-                    
-                    success, encoded_image = cv2.imencode('.jpg', image)
-                    if not success:
-                        logger.error('imencode image error')
-                        continue
-                    content = encoded_image.tobytes()
-                    successori, encoded_imageori = cv2.imencode('.jpg', frame)
-                    if not successori:
-                        logger.error('imencode original image error')
-                        continue
-                    contentori = encoded_imageori.tobytes()
-                    
-                    payload = {
-                        'channel': '45',
-                        'classIndex': '8',
-                        'ip': '172.19.152.231',
-                        'videoTime': videoTime,
-                        'videoUrl': stream.stream_url
-                    }
-                    
-                    files = [
-                        ('file', (filename, content, 'image/jpeg')),
-                        ('oldFile', (filenameori, contentori, 'image/jpeg'))
-                    ]
-                    
-                    # 使用协程上传图像
-                    await upload_image(session, 'http://172.19.152.231/open/api/operate/upload', payload, files)
-                except Exception as error:
-                    logger.error(f'Error: {str(error)}')
+            for frame, ret in stream():
+                if not is_time_in_period(start_det_time, end_det_time):
+                    logger.info(f"当前时间不在时间段 {start_det_time} ~ {end_det_time}")
+                    stream.close()
+                    break
+                if not ret: 
+                    continue
+                if time.time() - posttime < 30:
+                    continue
+                image = frame.copy()
+                result = instance(image)
+                if len(result) > 0:
+                    try:
+                        posttime = time.time()
+                        videoTime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
+                        fileTime = time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())
+                        filename = fileTime + ".jpg"
+                        filenameori = fileTime + "det.jpg"
+                        logger.info(videoTime)
+                        logger.info(result)
+                        
+                        for res in result:
+                            cv2.rectangle(image, tuple(map(int, (res.left, res.top))), 
+                                        tuple(map(int, (res.right, res.bottom))), (255, 0, 0), 4)
+                        
+                        success, encoded_image = cv2.imencode('.jpg', image)
+                        if not success:
+                            logger.error('imencode image error')
+                            continue
+                        content = encoded_image.tobytes()
+                        successori, encoded_imageori = cv2.imencode('.jpg', frame)
+                        if not successori:
+                            logger.error('imencode original image error')
+                            continue
+                        contentori = encoded_imageori.tobytes()
+                        
+                        payload = {
+                            'channel': '45',
+                            'classIndex': '8',
+                            'ip': '172.19.152.231',
+                            'videoTime': videoTime,
+                            'videoUrl': stream.stream_url
+                        }
+                        
+                        files = [
+                            ('file', (filename, content, 'image/jpeg')),
+                            ('oldFile', (filenameori, contentori, 'image/jpeg'))
+                        ]
+                        
+                        # 使用协程上传图像
+                        await upload_image(session, 'http://172.19.152.231/open/api/operate/upload', payload, files)
+                    except Exception as error:
+                        logger.error(f'Error: {str(error)}')
 
     logger.info("======= EXIT =======")
 

+ 3 - 0
stream.py

@@ -43,6 +43,9 @@ class StreamCapture:
         logger.info("Reconnect success")
         return True
 
+    def close(self):
+        self.cap.release()
+
     def __call__(self):
         stream_url = self.get_stream_url()
         logger.info(f"Connect {stream_url} ...")

+ 25 - 5
test.py

@@ -7,6 +7,21 @@ from logger import logger
 from stream import StreamCapture
 from infer import DoorInference
 from logger import logger
+from datetime import datetime
+
+def is_time_in_period(start_time_str, end_time_str):
+    # 解析开始和结束时间字符串为time对象
+    start_time = datetime.strptime(start_time_str, "%H:%M").time()
+    end_time = datetime.strptime(end_time_str, "%H:%M").time()
+    current_time = datetime.now().time()
+    
+    # 判断时间段是否跨天
+    if start_time <= end_time:
+        # 时间段在同一天
+        return start_time <= current_time <= end_time
+    else:
+        # 时间段跨天(如22:00到次日02:00)
+        return current_time >= start_time or current_time <= end_time
 
 async def upload_image(url, payload, filename1, content1, filename2, content2):
     form_data = aiohttp.FormData()
@@ -25,8 +40,13 @@ async def upload_image(url, payload, filename1, content1, filename2, content2):
             logger.info(result)
 
 async def process_stream():
+    start_det_time = "10:00"
+    end_det_time = "20:00"
+    if not is_time_in_period(start_det_time, end_det_time):
+        logger.info(f"当前时间不在时间段 {start_det_time} ~ {end_det_time}")
+        return
     logger.info("====== Start Server =======")
-    human_model_path = "models/work_clo_person_head_hat.pt"
+    human_model_path = "models/yolo11m.pt"
     door_model_path = "models/door_classify.pt"
     test_area = [[(222, 59), (432, 3), (528, 96), (318, 198)]]
     
@@ -34,7 +54,7 @@ async def process_stream():
     ip = '172.19.152.231'
     channel = '45'
     stream = StreamCapture(ip, channel)
-    stream.get_stream_url()
+    # stream.get_stream_url()
     
     posttime = time.time() - 30
     
@@ -43,7 +63,7 @@ async def process_stream():
     image = frame.copy()
         
     result = instance(image)
-    if len(result) > 0 and time.time() - posttime > 30:
+    if len(result) > 0 and time.time() - posttime >= 30:
         posttime = time.time()
         videoTime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
         fileTime = time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())
@@ -55,7 +75,7 @@ async def process_stream():
         for res in result:
             cv2.rectangle(image, tuple(map(int, (res.left, res.top))), 
                             tuple(map(int, (res.right, res.bottom))), (255, 0, 0), 4)
-        
+        cv2.imwrite("result.jpg", image)
         success, encoded_image = cv2.imencode('.jpg', image)
         if not success:
             logger.error('imencode image error')
@@ -74,7 +94,7 @@ async def process_stream():
         }
         
         # 使用协程上传图像
-        await upload_image('http://172.19.152.231/open/api/operate/upload', payload, filename, content, filenameori, contentori)
+        # await upload_image('http://172.19.152.231/open/api/operate/upload', payload, filename, content, filenameori, contentori)
 
     logger.info("======= EXIT =======")