Added video function to count frames and a test for it.

Author: fnunnari

PostProcessing/test_functions.py

@@ -11,6 +11,8 @@
 
 from PostProcessVideos import scan_session_dir
 
+from video import video_info
+
 
 RECSYNCH_SESSION_DIR_VAR = "RECSYNCH_SESSION_DIR"
 
@@ -97,3 +99,16 @@ def test_df_trimming(session_data):
         # Get the last element of the first column
         ts_end = df[0].iloc[-1]
         assert ts_end >= max_common
+
+
+@pytest.mark.parametrize("client_data", session_data_list())
+def test_video_sources(client_data):
+    _, _, video_filepath = client_data
+
+    video_path = Path(video_filepath)
+
+    assert video_path.exists()
+    assert video_path.is_file()
+
+    w, h, n_frames = video_info(video_path)
+    assert n_frames >= 2

PostProcessing/video.py

@@ -2,12 +2,44 @@
 import os 
 
 import pandas as pd
+import ffmpeg
+
+from typing import Tuple
+
+
+def video_info(video_path: str) -> Tuple[int, int, int]:
+    """
+    Uses the ffmpeg.probe function to retrieve information about a video file.
+
+    :param video_path: Path to a valid video file
+    :return: A 3-tuple with integers for (width, height, number_of_frames)
+    """
+
+    #
+    # Fetch video info
+    info = ffmpeg.probe(video_path)
+    # Get the list of all video streams
+    video_streams = [stream for stream in info['streams'] if stream['codec_type'] == 'video']
+    if len(video_streams) == 0:
+        raise BaseException("No video streams found in file '{}'".format(video_path))
+
+    # retrieve the first stream of type 'video'
+    info_video = video_streams[0]
+
+    video_w = info_video['width']
+    video_h = info_video['height']
+    n_frames = int(info_video['nb_frames'])
+
+    return video_w, video_h, n_frames
+
 
 def extract_frames(video_file: str, timestamps: pd.DataFrame, output_dir: str):
+
     # Open the video file
     cap = cv2.VideoCapture(video_file)
     assert cap.isOpened() == True
-    video_name = os.path.splitext(video_file)[0]
+    video_name, _ = os.path.splitext(video_file)
+
     # Loop over each timestamp in the CSV file
     for timestamp in timestamps:
         # Extract the frame using OpenCV
@@ -20,4 +52,3 @@ def extract_frames(video_file: str, timestamps: pd.DataFrame, output_dir: str):
 
     # Release the video file
     cap.release()
-