initial commit

person_yolo_server/person_yolo_infer_server.py

@@ -0,0 +1,53 @@
+import toml
+from loguru import logger
+import zmq
+import numpy as np
+import cv2
+
+from infer import Person_model_infer
+from visualize import visualize_box_mask
+
+labels = [
+    "pedestrian"
+]
+
+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 = Person_model_infer()
+    logger.info("person yolo model load success")
+
+    context1 = zmq.Context()
+    camera1_socket = context1.socket(zmq.REQ)
+    camera1_socket.connect(f"tcp://localhost:{cfg['camera']['camera1_port']}")
+    logger.info("connect camera1 success")
+
+    context2 = zmq.Context()
+    socket = context2.socket(zmq.REP)
+    socket.bind("tcp://*:7778")
+    logger.info("bind server success")
+    while True:
+        message = socket.recv()
+        camera1_socket.send_string("")
+        message = camera1_socket.recv()
+        np_array = np.frombuffer(message, dtype=np.uint8)
+        frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
+        result = predictor.infer(frame)
+        img = visualize_box_mask(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR),result,labels)
+        showim = np.array(img)
+        socket.send_pyobj(showim)
+        if cv2.waitKey(1) == 27:
+            break
+    camera1_socket.close()
+    socket.close()
+    context1.term()
+    context2.term()
+
+    logger.info("Interrupt received, stopping...")
+    logger.info("person yolo infer server exit")
+                
+