Преглед на файлове

git merge is so stupid, revert.

lyc8503 преди 4 години
родител
ревизия
b51bee1a4f
променени са 1 файла, в които са добавени 48 реда и са изтрити 48 реда
  1. 48 48
      backend_refactor/service/video_analyzer.py

+ 48 - 48
backend_refactor/service/video_analyzer.py

@@ -1,3 +1,4 @@
+import math
 import os
 from typing import Callable
 from config import UPLOAD_DIR
@@ -35,6 +36,37 @@ def split_video(video_uuid):
     cap.release()
 
 
+def gather_video(video_uuid):
+    video_path = UPLOAD_DIR + video_uuid + ".mp4"
+    frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
+
+    # Just to get a fps
+    cap = cv2.VideoCapture(video_path)
+    fps = cap.get(5)
+    cap.release()
+
+    frames = [name for name in os.listdir(frames_dir) if
+              name.endswith(".jpg") and "npy" not in name and "result" not in name]
+    num_frames = len(frames)
+
+    img_array = []
+    for i in range(1, num_frames + 1):
+        frame_path = frames_dir + str(i) + ".jpg"
+        img = cv2.imread(frame_path)
+        if img is None:
+            continue
+        img_array.append(img)
+
+    logging.debug("Frames total: %d" % (len(img_array)))
+
+    result_path = UPLOAD_DIR + video_uuid + ".result.mp4"
+
+    out = cv2.VideoWriter(result_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (img_array[0].shape[1], img_array[0].shape[1]))  # So fucking stupid
+    for i in img_array:
+        out.write(i)
+    out.release()
+
+
 from service.aiutil.torch_openpose import torch_openpose
 from service.aiutil.util import draw_bodypose
 
@@ -46,6 +78,7 @@ def run_openpose_for_frames(video_uuid):
     for i in os.listdir(frames_dir):
         if i.endswith(".jpg") and "result" not in i:
 
+            # Already processed
             if os.path.exists(frames_dir + i.replace(".jpg", ".result.jpg")) and os.path.exists(
                     frames_dir + i + ".npy"):
                 continue
@@ -57,20 +90,23 @@ def run_openpose_for_frames(video_uuid):
             np.save(frames_dir + i + ".npy", poses)
 
             canvas = draw_bodypose(oriImg, poses, 'body_25')
-
             cv2.imwrite(frames_dir + i.replace(".jpg", ".result.jpg"), canvas)
 
 
 class BaseVideoAnalyzer:
-    def __init__(self, video_uuid):
-        self.video_uuid = video_uuid
+    def __init__(self, data: dict):
+        self.video_uuid = data['video_uuid']
+
+    def _do_analyze(self):
+        assert False, "Internal error: You should override this with a subclass!"
 
     def analyze(self, callback: Callable):
         try:
             callback("RUNNING")
             split_video(self.video_uuid)
             run_openpose_for_frames(self.video_uuid)
-            callback("FINISHED", result={})
+            gather_video(self.video_uuid)
+            callback("FINISHED", result=self._do_analyze())
         except Exception as e:
             callback("ERROR", error=e)
             raise e
@@ -364,31 +400,13 @@ def ignore_data_jump(result_ig: dict):
 
 
 class HighKneesVideoAnalyzer(BaseVideoAnalyzer):
-    def __init__(self, video_uuid):
-        super().__init__(video_uuid)
-
-    def analyze(self, callback: Callable):
-        try:
-            callback("RUNNING")
-
-            callback("FINISHED", result={})
-        except Exception as e:
-            callback("ERROR", error=e)
-            raise e
+    def _do_analyze(self):
+        pass
 
 
 class CrunchVideoAnalyzer(BaseVideoAnalyzer):
-    def __init__(self, video_uuid):
-        super().__init__(video_uuid)
-
-    def analyze(self, callback: Callable):
-        try:
-            callback("RUNNING")
-
-            callback("FINISHED", result={})
-        except Exception as e:
-            callback("ERROR", error=e)
-            raise e
+    def _do_analyze(self):
+        pass
 
 
 class StandingLongJumpVideoAnalyzer(BaseVideoAnalyzer):
@@ -404,28 +422,10 @@ class StandingLongJumpVideoAnalyzer(BaseVideoAnalyzer):
 
 
 class PlankVideoAnalyzer(BaseVideoAnalyzer):
-    def __init__(self, video_uuid):
-        super().__init__(video_uuid)
-
-    def analyze(self, callback: Callable):
-        try:
-            callback("RUNNING")
-
-            callback("FINISHED", result={})
-        except Exception as e:
-            callback("ERROR", error=e)
-            raise e
+    def _do_analyze(self):
+        pass
 
 
 class SquatVideoAnalyzer(BaseVideoAnalyzer):
-    def __init__(self, video_uuid):
-        super().__init__(video_uuid)
-
-    def analyze(self, callback: Callable):
-        try:
-            callback("RUNNING")
-
-            callback("FINISHED", result={})
-        except Exception as e:
-            callback("ERROR", error=e)
-            raise e
+    def _do_analyze(self):
+        pass