Merge pull request #328 from fourMs/copilot/fix-optical-flow-export-issue

Fix optical flow video export failing with FFprobeError on frame count

Author: alexarje

musicalgestures/_flow.py

@@ -7,7 +7,7 @@
 from scipy.stats import entropy
 
 import musicalgestures
-from musicalgestures._utils import MgFigure, extract_wav, embed_audio_in_video, MgProgressbar, convert_to_avi, generate_outfilename
+from musicalgestures._utils import MgFigure, extract_wav, embed_audio_in_video, MgProgressbar, convert_to_avi, generate_outfilename, ffmpeg_cmd
 
 
 class Flow:
@@ -111,8 +111,10 @@ def dense(
             if not overwrite:
                 target_name = generate_outfilename(target_name)
 
-            fourcc = cv2.VideoWriter_fourcc(*'MJPG')
-            out = cv2.VideoWriter(target_name, fourcc, fps, (width, height))
+            cmd = ['ffmpeg', '-y', '-s', '{}x{}'.format(width, height),
+                   '-r', str(fps), '-f', 'rawvideo', '-pix_fmt', 'bgr24', '-vcodec', 'rawvideo',
+                   '-i', '-', '-vcodec', 'mjpeg', '-q:v', '3', target_name]
+            out = ffmpeg_cmd(cmd, total_time=length, pipe='write')
 
         ret, frame1 = vidcap.read()
         prev_frame = cv2.cvtColor(cv2.resize(frame1, size), cv2.COLOR_BGR2GRAY)
@@ -154,14 +156,14 @@ def dense(
 
                     if skip_empty:
                         if np.sum(rgb) > 0:
-                            out.write(rgb.astype(np.uint8))
+                            out.stdin.write(rgb.astype(np.uint8))
                         else:
                             if ii == 0:
-                                out.write(rgb.astype(np.uint8))
+                                out.stdin.write(rgb.astype(np.uint8))
                             else:
-                                out.write(prev_rgb.astype(np.uint8))
+                                out.stdin.write(prev_rgb.astype(np.uint8))
                     else:
-                        out.write(rgb.astype(np.uint8))
+                        out.stdin.write(rgb.astype(np.uint8))
 
                     if skip_empty:
                         if np.sum(rgb) > 0 or ii == 0:
@@ -232,7 +234,8 @@ def dense(
             return mgf
 
         else:
-            out.release()
+            out.stdin.close()
+            out.wait()
             destination_video = target_name
 
             if self.has_audio:
@@ -321,7 +324,6 @@ def sparse(
 
         vidcap = cv2.VideoCapture(filename)
         ret, frame = vidcap.read()
-        fourcc = cv2.VideoWriter_fourcc(*'MJPG')
 
         fps = int(vidcap.get(cv2.CAP_PROP_FPS))
         width = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
@@ -336,7 +338,10 @@ def sparse(
         if not overwrite:
             target_name = generate_outfilename(target_name)
 
-        out = cv2.VideoWriter(target_name, fourcc, fps, (width, height))
+        cmd = ['ffmpeg', '-y', '-s', '{}x{}'.format(width, height),
+               '-r', str(fps), '-f', 'rawvideo', '-pix_fmt', 'bgr24', '-vcodec', 'rawvideo',
+               '-i', '-', '-vcodec', 'mjpeg', '-q:v', '3', target_name]
+        out = ffmpeg_cmd(cmd, total_time=length, pipe='write')
 
         # params for ShiTomasi corner detection
         feature_params = dict(maxCorners=corner_max_corners,
@@ -390,7 +395,7 @@ def sparse(
 
                 img = cv2.add(frame, mask)
 
-                out.write(img.astype(np.uint8))
+                out.stdin.write(img.astype(np.uint8))
 
                 # Now update the previous frame and previous points
                 old_gray = frame_gray.copy()
@@ -403,7 +408,8 @@ def sparse(
             pb.progress(ii)
             ii += 1
 
-        out.release()
+        out.stdin.close()
+        out.wait()
 
         destination_video = target_name
 

tests/test_flow.py

@@ -0,0 +1,82 @@
+import musicalgestures
+import os
+import pytest
+
+
+@pytest.fixture(scope="class")
+def testvideo_avi(tmp_path_factory):
+    target_name = str(tmp_path_factory.mktemp("data")).replace("\\", "/") + "/testvideo.avi"
+    testvideo_avi = musicalgestures._utils.extract_subclip(
+        musicalgestures.examples.dance, 5, 6, target_name=target_name)
+    return testvideo_avi
+
+
+@pytest.fixture(scope="class")
+def testvideo_mp4(tmp_path_factory):
+    target_name = str(tmp_path_factory.mktemp("data")).replace(
+        "\\", "/") + "/testvideo.avi"
+    testvideo_avi = musicalgestures._utils.extract_subclip(
+        musicalgestures.examples.dance, 5, 6, target_name=target_name)
+    testvideo_mp4 = musicalgestures._utils.convert_to_mp4(testvideo_avi)
+    os.remove(testvideo_avi)
+    return testvideo_mp4
+
+
+class Test_flow_dense:
+    def test_normal_case(self, testvideo_avi):
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.dense()
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_overwrite(self, testvideo_avi):
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.dense(overwrite=True)
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_not_avi(self, testvideo_mp4):
+        mg = musicalgestures.MgVideo(testvideo_mp4)
+        result = mg.flow.dense()
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_skip_empty(self, testvideo_avi):
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.dense(skip_empty=True)
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_with_target_name(self, testvideo_avi):
+        target_name = os.path.dirname(testvideo_avi) + "/result_dense.avi"
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.dense(target_name=target_name, overwrite=True)
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+
+class Test_flow_sparse:
+    def test_normal_case(self, testvideo_avi):
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.sparse()
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_overwrite(self, testvideo_avi):
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.sparse(overwrite=True)
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_not_avi(self, testvideo_mp4):
+        mg = musicalgestures.MgVideo(testvideo_mp4)
+        result = mg.flow.sparse()
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True
+
+    def test_with_target_name(self, testvideo_avi):
+        target_name = os.path.dirname(testvideo_avi) + "/result_sparse.avi"
+        mg = musicalgestures.MgVideo(testvideo_avi)
+        result = mg.flow.sparse(target_name=target_name, overwrite=True)
+        assert type(result) == musicalgestures.MgVideo
+        assert os.path.isfile(result.filename) == True