Changed to ffmpeg

.gitignore

@@ -3,3 +3,4 @@ __pycache__/
 dataset.tar
 *.json
 *.mp4
+video/*.jpg

create_video.py

@@ -1,6 +1,7 @@
 import numpy as np
 import cv2
 from PIL import Image
+import ffmpeg
 
 
 def noise_process(numpy_image, steps=149):
@@ -29,10 +30,15 @@ def generate_video(numpy_image):
     image_lst = np.repeat(image_lst, copies, axis=0)
     image_lst = np.concatenate((image_lst, image_lst[:spill_over]), axis=0)
     image_lst = image_lst[::-1]
-    fourcc = cv2.VideoWriter_fourcc(*'avc1')
-    video = cv2.VideoWriter(save_path, fourcc, float(fps), (width, height))
-    for frame_count in range(fps * sec):
-        img = np.expand_dims(image_lst[frame_count], 2)
-        video.write(img.astype(np.uint8))
-    video.release()
+    for i, img in enumerate(image_lst):
+        Image.fromarray(img).save(f"video/{i}.jpg")
+
+    ffmpeg.input('video/*.jpg', pattern_type='glob', framerate=25).output(save_path).run()
+
+    # fourcc = cv2.VideoWriter_fourcc(*'avc1')
+    # video = cv2.VideoWriter(save_path, fourcc, float(fps), (width, height))
+    # for frame_count in range(fps * sec):
+    #     img = np.expand_dims(image_lst[frame_count], 2)
+    #     video.write(img.astype(np.uint8))
+    # video.release()
     return save_path

requirements.txt

@@ -5,6 +5,7 @@ pandas
 torch
 ffprobe
 pytorch-gradcam
+ffmpeg-python
 Pillow
 tqdm
 opencv-python