فهرست منبع

add picture return

lyc8503 4 سال پیش
والد
کامیت
8c689fd733

+ 14 - 2
backend_refactor/app.py

@@ -1,12 +1,13 @@
 import logging
+import os
 
-from flask import Flask
+from flask import Flask, send_from_directory
 from flask_cors import CORS
 
 import route.analyze
 import route.result
 import route.upload
-from config import VERSION_PREFIX, MAX_CONTENT_LENGTH, MYSQL_URI
+from config import VERSION_PREFIX, MAX_CONTENT_LENGTH, MYSQL_URI, UPLOAD_DIR
 from db import db
 from route.util import make_response
 
@@ -29,11 +30,22 @@ with app.app_context():
     db.create_all()
 
 
+try:
+    os.mkdir(UPLOAD_DIR + "/result")
+except:
+    pass
+
+
 @app.route("/")
 def hello_world():
     return "<p>Hello, World!</p>"
 
 
+@app.route("/file/<path:name>")
+def download_file(name):
+    return send_from_directory(UPLOAD_DIR + '/result/', name)
+
+
 @app.errorhandler(413)
 def error_handler(e):
     logging.exception(e)

+ 1 - 0
backend_refactor/config/__init__.py

@@ -3,5 +3,6 @@ import os
 VERSION_PREFIX = "/v1/"
 UPLOAD_DIR = "/data/"
 MAX_CONTENT_LENGTH = 100 * 1024 * 1024
+FILE_URL = "http://106.15.1.178:8090/file/"
 
 MYSQL_URI = f'mysql://{os.environ["MYSQL_USER"]}:{os.environ["MYSQL_PASSWORD"]}@{os.environ["MYSQL_ADDRESS"]}/{os.environ["MYSQL_DATABASE"]}'

+ 6 - 5
backend_refactor/service/photo_analyzer.py

@@ -3,7 +3,7 @@ from typing import Callable
 
 import cv2
 
-from config import UPLOAD_DIR
+from config import UPLOAD_DIR, FILE_URL
 from service.ai.photo.standing_photo import *
 from service.ai.util.torch_openpose import torch_openpose
 from service.ai.util.util import draw_bodypose
@@ -25,7 +25,7 @@ class StandingPhotoAnalyzer(BasePhotoAnalyzer):
 
     def analyze(self, callback: Callable) -> None:
         try:
-            callback("RUNNING",progress=0)
+            callback("RUNNING", progress=0)
             photos = []
             npys = []
 
@@ -45,10 +45,11 @@ class StandingPhotoAnalyzer(BasePhotoAnalyzer):
                 canvas = copy.deepcopy(ori_img)
                 canvas = draw_bodypose(canvas, poses, 'body_25')
 
-                cv2.imwrite(UPLOAD_DIR + '{}.result.jpg'.format(i), canvas)
-                photos.append(UPLOAD_DIR + '{}.result.jpg'.format(i))
+                cv2.imwrite(UPLOAD_DIR + 'result/{}.result.jpg'.format(i), canvas)
+                photos.append(FILE_URL + '{}.result.jpg'.format(i))
 
-            callback("FINISHED", result={"data": {}, "advice": {"front": npys[0], "right": npys[1]}})
+            callback("FINISHED", result={"data": {"front": {"file": photos[0]}, "right": {"file": photos[1]}},
+                                         "advice": {"front": npys[0], "right": npys[1]}})
         except Exception as e:
             callback("ERROR", error=e)
             raise e

+ 12 - 3
backend_refactor/service/video_analyzer.py

@@ -5,7 +5,7 @@ from typing import Callable
 
 import cv2
 
-from config import UPLOAD_DIR
+from config import UPLOAD_DIR, FILE_URL
 from service.ai.util.torch_openpose import torch_openpose
 from service.ai.util.util import draw_bodypose
 from service.ai.video.classify.balancebeam import classify_balancebeam
@@ -61,9 +61,10 @@ def run_openpose_for_npy_and_video(video_uuid, callback: Callable):
                 results[str(c)] = poses  # this is evil, but the AI code expect it to be so.
 
                 canvas = draw_bodypose(frame, poses, 'body_25')
+                cv2.imwrite(UPLOAD_DIR + "result/" + video_uuid + "." + str(c) + ".jpg", canvas)
 
                 if out is None:
-                    out = cv2.VideoWriter(UPLOAD_DIR + video_uuid + ".result.mp4", cv2.VideoWriter_fourcc(*"mp4v"), FPS,
+                    out = cv2.VideoWriter(UPLOAD_DIR + "result/" + video_uuid + ".result.mp4", cv2.VideoWriter_fourcc(*"mp4v"), FPS,
                                           (canvas.shape[1], canvas.shape[0]))  # So fucking stupid
                 out.write(canvas)
 
@@ -98,7 +99,15 @@ class BaseVideoAnalyzer:
             else:  # read cache
                 self.results = pickle.load(open(PICKLE_DUMP_PATH, "rb"))
 
-            callback("FINISHED", result=self._do_analyze())
+            result = {**self._do_analyze(), "file": FILE_URL + self.video_uuid + ".result.mp4"}
+            # inject file url here
+            try:
+                for i in result['data']:
+                    result['data'][i]['file'] = FILE_URL + self.video_uuid + "." + result['data'][i]['index'] + ".jpg"
+            except Exception as e:
+                logging.warning("file index in data not found, %s" % e)
+
+            callback("FINISHED", result=result)
         except Exception as e:
             callback("ERROR", error=e)
             raise e