diff --git a/server.py b/server.py index 13e578e..5366f29 100644 --- a/server.py +++ b/server.py @@ -207,6 +207,13 @@ def parse_args(): return opt +import sys +import signal + +def quit(signum, frame): + print('stop fusion') + sys.exit() + def put_watermark(img, wm_encoder=None): if wm_encoder is not None: img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR) @@ -214,7 +221,6 @@ def put_watermark(img, wm_encoder=None): img = Image.fromarray(img[:, :, ::-1]) return img - # 对任务状态的修改 def update_task_status(task: dict, status: str, progress: int, data: list = []): task["status"] = status @@ -224,6 +230,8 @@ def update_task_status(task: dict, status: str, progress: int, data: list = []): # 从局域网中获取一组任务(如果列表为空,等待2s后重新获取) def get_tasks(tasks:list=[]): + signal.signal(signal.SIGINT, quit) + signal.signal(signal.SIGTERM, quit) while len(tasks) == 0: try: tasks = requests.get("http://localhost:3000/api/drawing?status=waiting").json() @@ -234,12 +242,6 @@ def get_tasks(tasks:list=[]): time.sleep(2) return tasks -import sys -import signal - -def quit(signum, frame): - print('stop fusion') - sys.exit() def main_dev(opt): signal.signal(signal.SIGINT, quit)