From 90f5f92abbd53b11ea63d4ef02a1aa5c65827e66 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=A1=9C=E8=8F=AF?= <huan0016@gmail.com>
Date: Wed, 1 Mar 2023 02:04:56 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E9=9B=85=E9=80=80=E5=87=BA2?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

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)