Parcourir la source

Add task progress for backend.

Bob Huang il y a 3 ans
Parent
commit
6fd5c9aec4

+ 3 - 1
backend_refactor/db/task.py

@@ -12,11 +12,13 @@ class Task(db.Model):
                              'VIDEO_SINGLELEGSTAND', 'VIDEO_PLANK', 'VIDEO_SQUAT'))
     create_time = db.Column(db.BigInteger)
     status = db.Column(db.Enum('QUEUEING', 'RUNNING', 'FINISHED', 'ERROR'))
+    progress = db.Column(db.Integer)
     result = db.Column(db.String(8192))
 
-    def __init__(self, uuid, type, create_time=time.time() * 1000, status='QUEUEING', result='{}'):
+    def __init__(self, uuid, type, create_time=time.time() * 1000, status='QUEUEING', progress=0, result='{}'):
         self.uuid = uuid
         self.type = type
         self.create_time = create_time
         self.status = status
+        self.progress = progress
         self.result = result

+ 2 - 0
backend_refactor/route/analyze.py

@@ -28,6 +28,8 @@ def callback_generator(task_uuid):
             with app.app_context():
                 task = Task.query.get(task_uuid)
                 task.status = status
+                if status == 'RUNNING':
+                    task.progress = kwargs['progress']
                 if status == 'FINISHED':
                     task.result = json.dumps(kwargs['result'])
                 if status == 'ERROR':

+ 4 - 4
backend_refactor/route/result.py

@@ -20,12 +20,12 @@ def get_result(task_uuid):
 
     if status == 'FINISHED':
         return make_response(200, "success",
-                             {"task_uuid": task.uuid, "type": task.type, "result": json.loads(task.result)})
+                             {"task_uuid": task.uuid, "type": task.type, "status": "FINISHED", "result": json.loads(task.result)})
     elif status == 'QUEUEING':
-        return make_response(200, "queueing", {"task_uuid": task.uuid, "type": task.type})
+        return make_response(200, "success", {"task_uuid": task.uuid, "type": task.type, "status": "QUEUEING"})
     elif status == 'RUNNING':
-        return make_response(200, "running", {"task_uuid": task.uuid, "type": task.type})
+        return make_response(200, "success", {"task_uuid": task.uuid, "type": task.type, "status": "RUNNING", "progress": task.progress})
     elif status == 'ERROR':
-        return make_response(200, "error", {"task_uuid": task.uuid, "type": task.type})
+        return make_response(200, "success", {"task_uuid": task.uuid, "type": task.type, "status": "ERROR"})
     else:
         assert 0.1 + 0.2 == 0.3, "what do u mean?"

+ 5 - 3
backend_refactor/service/video_analyzer.py

@@ -34,11 +34,12 @@ from service.ai.video.video_score.tennisthrow import cal_tennisthrow
 from service.ai.video.video_score.turnaround import cal_turnaround
 
 
-def run_openpose_for_npy_and_video(video_uuid):
+def run_openpose_for_npy_and_video(video_uuid, callback: Callable):
     video_path = UPLOAD_DIR + video_uuid + ".mp4"
 
     cap = cv2.VideoCapture(video_path)
-    FPS = cap.get(5)  # 5 means fps
+    FPS = cap.get(cv2.CAP_PROP_FPS)
+    total_frames = cap.get(cv2.CAP_PROP_FRAME_COUNT)
     out = None
     logging.debug("Splitting video %s, FPS: %s" % (video_path, FPS))
 
@@ -67,6 +68,7 @@ def run_openpose_for_npy_and_video(video_uuid):
                 out.write(canvas)
 
             c += 1
+            callback("RUNNING", progress=int(100*c/(total_frames/sampling_fps)))
         else:
             break
 
@@ -91,7 +93,7 @@ class BaseVideoAnalyzer:
             PICKLE_DUMP_PATH = UPLOAD_DIR + self.video_uuid + ".dump"
 
             if not os.path.exists(PICKLE_DUMP_PATH):  # analyze
-                self.results = run_openpose_for_npy_and_video(self.video_uuid)
+                self.results = run_openpose_for_npy_and_video(self.video_uuid, callback)
                 pickle.dump(self.results, open(PICKLE_DUMP_PATH, "wb"))
             else:  # read cache
                 self.results = pickle.load(open(PICKLE_DUMP_PATH, "rb"))