54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
|
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")
|
||
|
|
|
||
|
|
|