From 042f6ece691a545d8931c4ce797edd06a01e0d29 Mon Sep 17 00:00:00 2001
From: 2024snow <2103200855@qq.com>
Date: Wed, 3 Jul 2024 18:27:18 +0800
Subject: [PATCH] update yolo_infer_server.py

---
 yolo_server/yolo_infer_server.py | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/yolo_server/yolo_infer_server.py b/yolo_server/yolo_infer_server.py
index a521bdc..330327e 100644
--- a/yolo_server/yolo_infer_server.py
+++ b/yolo_server/yolo_infer_server.py
@@ -1,6 +1,7 @@
 import toml
 import threading
 from loguru import logger
+import logging
 import zmq
 import numpy as np
 import cv2
@@ -22,10 +23,7 @@ start = False
 exit_event = threading.Event()
 
 
-context2 = zmq.Context()
-socket_server = context2.socket(zmq.PUB)
-socket_server.setsockopt(zmq.SNDHWM,10)
-socket_server.bind("tcp://*:7777")
+
 
 
 labels = [
@@ -100,8 +98,24 @@ def camera_resp(camera1_port, camera2_port):
     context.term()
     context1.term()
 
-if __name__ == "__main__":
-    cfg = toml.load('../cfg_infer_server.toml')
+
+def yolo_infer_server_main(queue):
+    if queue != None:
+        class Handler(logging.Handler):
+            def emit(self, record):
+                log_entry = self.format(record)
+                queue.put({'level': record.levelname.lower(), 'content': log_entry})
+        # logger.remove()
+        handler = Handler()
+        logger.add(handler, format="{time:MM-DD HH:mm:ss} {message}", level="DEBUG")
+    
+    context2 = zmq.Context()
+    socket_server = context2.socket(zmq.PUB)
+    socket_server.setsockopt(zmq.SNDHWM,10)
+    socket_server.bind("tcp://*:7777")
+
+
+    cfg = toml.load('/home/evan/Workplace/project_infer/cfg_infer_server.toml')
 
     # Configure log output
     logger.add(cfg['debug']['logger_filename'], format=cfg['debug']['logger_format'], retention=5, level="INFO")
@@ -149,3 +163,5 @@ if __name__ == "__main__":
     #     mythread1.join()
     #     mythread2.join()
     #     logger.info("yolo infer server exit")
+if __name__ == "__main__":
+    yolo_infer_server_main(None)
\ No newline at end of file