refactor gather video into one func.

backend_refactor/service/photo_analyzer.py

@@ -371,8 +371,8 @@ class StandingPhotoAnalyzer(BasePhotoAnalyzer):
                 canvas = copy.deepcopy(ori_img)
                 canvas = aiutil.util.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.jpg'.format(i), canvas)
+                photos.append(UPLOAD_DIR + '{}.result.jpg'.format(i))
 
             callback("FINISHED", result={"photos": {"front": photos[0], "right": photos[1]},
                                          "text": {"front": npys[0], "right": npys[1]}})

backend_refactor/service/video_analyzer.py

@@ -6,8 +6,11 @@ import cv2
 import logging
 import numpy as np
 
+from service.aiutil.torch_openpose import torch_openpose
+from service.aiutil.util import draw_bodypose
 
-def split_video(video_uuid):
+
+def run_openpose_for_npy_and_video(video_uuid):
     video_path = UPLOAD_DIR + video_uuid + ".mp4"
     frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
 
@@ -16,83 +19,41 @@ def split_video(video_uuid):
 
     cap = cv2.VideoCapture(video_path)
     FPS = cap.get(5)  # 5 means fps
+    out = None
     logging.debug("Splitting video %s, FPS: %s" % (video_path, FPS))
 
     c = 1
     sampling_fps = 1
 
+    tp = torch_openpose('body_25')  # AI Model
+
     while True:
         ret, frame = cap.read()
         if ret:
             if c % sampling_fps == 0:
-                target_path = frames_dir + str(c) + ".jpg"
-                if not os.path.exists(target_path):
-                    cv2.imwrite(target_path, frame)
-                    logging.debug("Writing frame %s, %s" % (c, target_path))
-            c += 1
-        else:
-            break
 
-    cap.release()
+                # if os.path.exists(frames_dir + str(c) + ".npy"):
+                #     continue
 
+                logging.debug("Running openpose for %s %s" % (video_uuid, i))
 
-def gather_video(video_uuid):
-    video_path = UPLOAD_DIR + video_uuid + ".mp4"
-    frames_dir = UPLOAD_DIR + "frames/" + video_uuid + "/"
+                poses = tp(frame)
+                np.save(frames_dir + str(c) + ".npy", poses)
 
-    # Just to get a fps
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(5)
-    cap.release()
+                canvas = draw_bodypose(frame, poses, 'body_25')
 
-    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)
+                if out is None:
+                    out = cv2.VideoWriter(frames_dir + "result.mp4", cv2.VideoWriter_fourcc(*"mp4v"), FPS, (canvas.shape[1], canvas.shape[0]))  # So fucking stupid
+                out.write(canvas)
 
-    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"
+            c += 1
+        else:
+            break
 
-    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)
+    cap.release()
     out.release()
 
 
-from service.aiutil.torch_openpose import torch_openpose
-from service.aiutil.util import draw_bodypose
-
-
-def run_openpose_for_frames(video_uuid):
-    tp = torch_openpose('body_25')
-
-    frames_dir = UPLOAD_DIR + "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
-
-            logging.debug("Running openpose for %s %s" % (video_uuid, i))
-
-            oriImg = cv2.imread(frames_dir + i)
-            poses = tp(oriImg)
-            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, data: dict):
         self.video_uuid = data['video_uuid']
@@ -103,9 +64,7 @@ class BaseVideoAnalyzer:
     def analyze(self, callback: Callable):
         try:
             callback("RUNNING")
-            split_video(self.video_uuid)
-            run_openpose_for_frames(self.video_uuid)
-            gather_video(self.video_uuid)
+            run_openpose_for_npy_and_video(self.video_uuid)
             callback("FINISHED", result=self._do_analyze())
         except Exception as e:
             callback("ERROR", error=e)