LeoDu 4 年之前
父节点
当前提交
22498d4585
共有 1 个文件被更改,包括 3 次插入13 次删除
  1. 3 13
      backend_refactor/service/video_analyzer.py

+ 3 - 13
backend_refactor/service/video_analyzer.py

@@ -45,23 +45,13 @@ def gather_video(video_uuid):
     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)
-
+    frames = list(sorted([name for name in os.listdir(frames_dir) if name.endswith(".result.jpg")], key=lambda x:int(x.replace(".result.jpg", ""))))
+    img_array = [cv2.imread(frames_dir + i) for i in frames]
     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
+    out = cv2.VideoWriter(result_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (img_array[0].shape[1], img_array[0].shape[0]))  # So fucking stupid
     for i in img_array:
         out.write(i)
     out.release()