فهرست منبع

抠图接口增加返回路径

tuon 1 سال پیش
والد
کامیت
b9947d019f
3فایلهای تغییر یافته به همراه27 افزوده شده و 6 حذف شده
  1. 20 0
      api/resp.py
  2. 4 3
      main.py
  3. 3 3
      ppmatting/core/predict.py

+ 20 - 0
api/resp.py

@@ -0,0 +1,20 @@
+class ApiResp:
+    # 基本响应
+    code = 0
+    msg = "Successful"
+    data = None
+
+    def __init__(self, code, msg=None, data=None):
+        self.code = code
+        if msg is not None and len(msg) > 0:
+            self.msg = msg
+        if data is not None:
+            self.data = data
+
+
+def error_resp(msg, code=-1):
+    return ApiResp(code=code, msg=msg)
+
+
+def success_resp(data):
+    return ApiResp(code=0, data=data)

+ 4 - 3
main.py

@@ -1,7 +1,8 @@
-from flask import Flask, request
+from flask import Flask, request, jsonify
 import os
 import tools
 from werkzeug.utils import secure_filename
+from api import resp
 
 app = Flask(__name__)
 
@@ -29,8 +30,8 @@ def seg():
     file_path = get_upload_file_path(filename)
     file.save(file_path)
 
-    tools.seg(file_path, save_dir)
-    return '{"code": 1}'
+    _, _, _, path = tools.seg(file_path, save_dir)
+    return jsonify(resp.success_resp(path))
 
 
 if __name__ == '__main__':

+ 3 - 3
ppmatting/core/predict.py

@@ -73,7 +73,7 @@ def save_result(alpha, path, im_path, trimap=None, fg_estimate=True, fg=None):
     rgba = np.concatenate((fg, alpha), axis=-1)
     cv2.imwrite(rgba_save_path, rgba)
 
-    return fg
+    return fg, alpha_save_path, rgba_save_path
 
 
 def reverse_transform(img, trans_info):
@@ -188,7 +188,7 @@ def predict(model,
 
             save_path = os.path.join(save_dir, im_file)
             mkdir(save_path)
-            fg = save_result(
+            fg,alpha_save_path, rgba_save_path = save_result(
                 alpha,
                 save_path,
                 im_path=im_path,
@@ -214,4 +214,4 @@ def predict(model,
             preprocess_cost_averager.reset()
             infer_cost_averager.reset()
             postprocess_cost_averager.reset()
-    return alpha, fg
+    return alpha, fg, alpha_save_path, rgba_save_path