102 lines
3.2 KiB
Python
102 lines
3.2 KiB
Python
import toml
|
|
import threading
|
|
from loguru import logger
|
|
import zmq
|
|
from infer_new import Yolo_model_infer
|
|
import numpy as np
|
|
import cv2
|
|
import time
|
|
lock1 = threading.Lock()
|
|
lock2 = threading.Lock()
|
|
lock3 = threading.Lock()
|
|
src_camera_id = 1
|
|
response = {'code': 0, 'data': []}
|
|
frame = None
|
|
start = False
|
|
# 处理 server 响应数据
|
|
def server_resp(yolo_infer_port):
|
|
logger.info("yolo server thread init success")
|
|
global response
|
|
global src_camera_id
|
|
|
|
context = zmq.Context()
|
|
# 启动 server
|
|
socket = context.socket(zmq.REP)
|
|
socket.bind(f"tcp://*:{yolo_infer_port}")
|
|
logger.info("yolo infer server init success")
|
|
while True:
|
|
message = socket.recv_string()
|
|
# 发送字符 1 和 2 切换摄像头 空字符表示请求推理数据
|
|
if message != '':
|
|
with lock1:
|
|
src_camera_id = int(message)
|
|
socket.send_pyobj({'code': 0, 'data': []})
|
|
else:
|
|
with lock2:
|
|
socket.send_pyobj(response)
|
|
|
|
# 处理摄像头数据
|
|
def camera_resp(camera1_port, camera2_port):
|
|
global frame
|
|
global src_camera_id
|
|
global start
|
|
context = zmq.Context()
|
|
camera1_socket = context.socket(zmq.SUB)
|
|
camera1_socket.connect(f"tcp://localhost:{camera1_port}")
|
|
camera1_socket.setsockopt_string(zmq.SUBSCRIBE, "")
|
|
logger.info("connect camera1 success")
|
|
context1 = zmq.Context()
|
|
camera2_socket = context1.socket(zmq.SUB)
|
|
camera2_socket.connect(f"tcp://localhost:{camera2_port}")
|
|
camera2_socket.setsockopt_string(zmq.SUBSCRIBE, "")
|
|
logger.info("connect camera2 success")
|
|
|
|
while True:
|
|
logger.info('111')
|
|
with lock1:
|
|
if src_camera_id == 1:
|
|
message = camera1_socket.recv()
|
|
else:
|
|
message = camera2_socket.recv()
|
|
|
|
# logger.info('111')
|
|
np_array = np.frombuffer(message, dtype=np.uint8)
|
|
with lock3:
|
|
frame = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
|
|
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
|
start = True
|
|
|
|
if __name__ == "__main__":
|
|
|
|
cfg = toml.load('../cfg_infer_server.toml')
|
|
|
|
# 配置日志输出
|
|
logger.add(cfg['debug']['logger_filename'], format=cfg['debug']['logger_format'], retention = 5, level="INFO")
|
|
|
|
|
|
# 初始化 paddle 推理器
|
|
predictor = Yolo_model_infer()
|
|
logger.info("yolo model load success")
|
|
# 启动 线程 1
|
|
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']['camera2_port']
|
|
),
|
|
daemon=True)
|
|
mythread1.start()
|
|
mythread2.start()
|
|
while True:
|
|
with lock3:
|
|
if start:
|
|
result = predictor.infer(frame)
|
|
with lock2:
|
|
response['data'] = result
|
|
mythread1.join()
|
|
mythread2.join()
|
|
logger.info("yolo infer server exit")
|
|
|