152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
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")
|