Bläddra i källkod

add stacktrace support for API

Jean-Louis Queguiner 6 år sedan
förälder
incheckning
80442e4136
1 ändrade filer med 35 tillägg och 21 borttagningar
  1. 35 21
      app.py

+ 35 - 21
app.py

@@ -16,6 +16,7 @@ import torch
 import fastai
 from fasterai.visualize import *
 from pathlib import Path
+import traceback
 
 
 torch.backends.cudnn.benchmark=True
@@ -31,42 +32,55 @@ app = Flask(__name__)
 
 @app.route("/process_image", methods=["POST"])
 def process_image():
-    source_url = request.json["source_url"]
-    render_factor = int(request.json["render_factor"])
+    try:
+        source_url = request.json["source_url"]
+        render_factor = int(request.json["render_factor"])
 
-    upload_directory = 'upload'
-    if not os.path.exists(upload_directory):
-           os.mkdir(upload_directory)
+        upload_directory = 'upload'
+        if not os.path.exists(upload_directory):
+            os.mkdir(upload_directory)
 
-    random_filename = str(uuid4()) + '.png'
+        random_filename = str(uuid4()) + '.png'
     
-    image_colorizer.plot_transformed_image_from_url(url=source_url, path=os.path.join(upload_directory, random_filename), figsize=(20,20),
+        image_colorizer.plot_transformed_image_from_url(url=source_url, path=os.path.join(upload_directory, random_filename), figsize=(20,20),
             render_factor=render_factor, display_render_factor=True, compare=False)
 
-    callback = send_file(os.path.join("result_images", random_filename), mimetype='image/jpeg')
+        callback = send_file(os.path.join("result_images", random_filename), mimetype='image/jpeg')
+        
+        return callback
 
-    os.remove(os.path.join("result_images", random_filename))
-    os.remove(os.path.join("upload", random_filename))
+    except:
+        traceback.print_exc()
+        return {message: 'input error'}, 400
+
+    finally:
+        os.remove(os.path.join("result_images", random_filename))
+        os.remove(os.path.join("upload", random_filename))
 
-    return callback
 
 @app.route("/process_video", methods=["POST"])
 def process_video():
-    source_url = request.json["source_url"]
-    render_factor = int(request.json["render_factor"])
+    try:
+        source_url = request.json["source_url"]
+        render_factor = int(request.json["render_factor"])
+
+        upload_directory = 'upload'
+        if not os.path.exists(upload_directory):
+            os.mkdir(upload_directory)
 
-    upload_directory = 'upload'
-    if not os.path.exists(upload_directory):
-           os.mkdir(upload_directory)
+        random_filename = str(uuid4()) + '.mp4'
 
-    random_filename = str(uuid4()) + '.mp4'
+        video_path = video_colorizer.colorize_from_url(source_url, random_filename, render_factor)
+        callback = send_file(os.path.join("video/result/", random_filename), mimetype='application/octet-stream')
 
-    video_path = video_colorizer.colorize_from_url(source_url, random_filename, render_factor)
-    callback = send_file(os.path.join("video/result/", random_filename), mimetype='application/octet-stream')
+        return callback
 
-    os.remove(os.path.join("video/result/", random_filename))
+    except:
+        traceback.print_exc()
+        return {message: 'input error'}, 400
 
-    return callback
+    finally:
+        os.remove(os.path.join("video/result/", random_filename))
 
 if __name__ == '__main__':
     port = 5000