initial commit

yolo_server/yolo_infer_server.py

@@ -0,0 +1,151 @@
+import toml
+import threading
+from loguru import logger
+import zmq
+import numpy as np
+import cv2
+import time
+
+# Custom imports
+from infer_new import Yolo_model_infer
+from visualize import visualize_box_mask
+# Initialize locks
+lock1 = threading.Lock()
+lock2 = threading.Lock()
+lock3 = threading.Lock()
+
+# Global variables
+src_camera_id = 1
+response = {'code': 0, 'data': []}
+frame = None
+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 = [
+    "tplatform", "tower", "sign", "shelter", "hospital", "basket", "base", 
+    "Yball", "Spiller", "Rmark", "Rblock", "Rball", "Mpiller", 
+    "Lpiller", "Lmark", "Bblock", "Bball"
+]
+
+# Handle server response data
+def server_resp(yolo_infer_port):
+    logger.info("yolo server thread init success")
+    global response
+    global src_camera_id
+
+    context = zmq.Context()
+    # Start server
+    socket = context.socket(zmq.REP)
+    socket.bind(f"tcp://*:{yolo_infer_port}")
+    logger.info("yolo infer server init success")
+    while not exit_event.is_set():
+        try:
+            message = socket.recv_string()
+            # Send character 1 and 2 to switch camera, empty string requests inference data
+            if message != '':
+                with lock1:
+                    logger.error(message)
+                    src_camera_id = int(message)
+                    logger.info("switch camera")
+                    socket.send_pyobj(response) 
+            else:
+                with lock2:
+                    socket.send_pyobj(response)
+                    response['data'] = np.array([])
+        except zmq.Again:
+            time.sleep(0.01)
+
+    socket.close()
+    context.term()
+
+# Handle camera data
+def camera_resp(camera1_port, camera2_port):
+    global frame
+    global src_camera_id
+    global start
+    context = zmq.Context()
+    camera1_socket = context.socket(zmq.REQ)
+    camera1_socket.connect(f"tcp://localhost:{camera1_port}")
+    logger.info("connect camera1 success")
+    context1 = zmq.Context()
+    camera2_socket = context1.socket(zmq.REQ)
+    camera2_socket.connect(f"tcp://localhost:{camera2_port}")
+    logger.info("connect camera2 success")
+
+    while not exit_event.is_set():
+        with lock1:
+            try:
+                if src_camera_id == 1:
+                    camera1_socket.send_string("")
+                    message = camera1_socket.recv()
+                else:
+                    camera2_socket.send_string("")
+                    message = camera2_socket.recv()
+                np_array = np.frombuffer(message, dtype=np.uint8)
+                with lock3:
+                    frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
+                    start = True
+            except:
+                time.sleep(0.01)
+
+    camera1_socket.close()
+    camera2_socket.close()
+    context.term()
+    context1.term()
+
+if __name__ == "__main__":
+    cfg = toml.load('../cfg_infer_server.toml')
+
+    # Configure log output
+    logger.add(cfg['debug']['logger_filename'], format=cfg['debug']['logger_format'], retention=5, level="INFO")
+
+    # Initialize YOLO inference model
+    predictor = Yolo_model_infer()
+    logger.info("yolo model load success")
+
+    # Start threads
+    mythread1 = threading.Thread(target=server_resp, args=(cfg['server']['yolo_infer_port'],), daemon=True)
+    mythread2 = threading.Thread(target=camera_resp, args=(cfg['camera']['camera1_port'], cfg['camera']['front_camera_port']), daemon=True)
+    mythread1.start()
+    mythread2.start()
+    while not exit_event.is_set():
+        with lock3:
+            if start:
+                result = predictor.infer(frame)
+                img = visualize_box_mask(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR),result,labels)
+                showim = np.array(img)
+                socket_server.send_pyobj(showim)
+                with lock2:
+                    response['data'] = result
+        # time.sleep(0.01)
+        if cv2.waitKey(1) == 27:
+            break
+    logger.info("Interrupt received, stopping...")
+    exit_event.set()
+    mythread1.join()
+    mythread2.join()
+    logger.info("yolo infer server exit")
+    # try:
+    #     while not exit_event.is_set():
+    #         with lock3:
+    #             if start:
+    #                 result = predictor.infer(frame)
+    #                 img = visualize_box_mask(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR),result,labels)
+    #                 showim = np.array(img)
+    #                 socket_server.send_pyobj(showim)
+    #                 with lock2:
+    #                     response['data'] = result
+    #         time.sleep(0.01)
+    # except KeyboardInterrupt:
+    #     logger.info("Interrupt received, stopping...")
+    #     exit_event.set()
+    #     mythread1.join()
+    #     mythread2.join()
+    #     logger.info("yolo infer server exit")