Add end to end tests

.github/workflows/ci.yml

@@ -70,7 +70,7 @@ jobs:
         run: |
           docker run -v ${{ github.workspace }}:/home/video_cap \
             mv-extractor:local \
-              python3.10 tests/tests.py
+              python3.10 -m unittest discover -s tests -p "*tests.py"
 
   build_and_test_wheels:
     name: Build wheels for cp${{ matrix.python }}-${{ matrix.platform_id }}
@@ -133,7 +133,7 @@ jobs:
           CIBW_MANYLINUX_X86_64_IMAGE: ${{ matrix.manylinux_image }}
           #CIBW_MANYLINUX_I686_IMAGE: ${{ matrix.manylinux_image }}
           CIBW_BUILD_FRONTEND: build
-          CIBW_TEST_COMMAND: VIDEO_URL={project}/vid_h264.mp4 python3 {project}/tests/tests.py
+          CIBW_TEST_COMMAND: PROJECT_ROOT={project} python3 -m unittest discover -s {project}/tests -p "*tests.py"
           CIBW_BUILD_VERBOSITY: 1
 
       - uses: actions/upload-artifact@v4

README.md

@@ -66,6 +66,7 @@ For example, to store extracted frames and motion vectors to disk without showin
 ```
 extract_mvs vid_h264.mp4 --dump
 ```
+The `--dump` parameter also takes an optional destination directory.
 
 
 ## Advanced Usage
@@ -79,13 +80,13 @@ git clone https://github.com/LukasBommes/mv-extractor.git mv_extractor
 
 Now, to run the tests from the `mv_extractor` directory with
 ```
-python3 tests/tests.py
+python3 -m unittest discover -s tests -p "*tests.py"
 ```
 Confirm that all tests pass.
 
 If you are using the Docker image instead of the PyPI package as explained below, you can invoke the tests with
 ```
-sudo ./run.sh python3.10 tests/tests.py
+sudo ./run.sh python3.10 -m unittest discover -s tests -p "*tests.py"
 ```
 
 ### Importing mvextractor into Your Own Scripts

src/mvextractor/__main__.py

@@ -25,16 +25,20 @@ def main(args=None):
         args = sys.argv[1:]
 
     parser = argparse.ArgumentParser(description='Extract motion vectors from video.')
-    parser.add_argument('video_url', type=str, nargs='?', help='File path or url of the video stream')
-    parser.add_argument('-p', '--preview', action='store_true', help='Show a preview video with overlaid motion vectors')
-    parser.add_argument('-v', '--verbose', action='store_true', help='Show detailled text output')
-    parser.add_argument('-d', '--dump', action='store_true', help='Dump frames, motion vectors, frame types, and timestamps to output directory')
+    parser.add_argument('video_url', type=str, nargs='?', help='file path or url of the video stream')
+    parser.add_argument('-p', '--preview', action='store_true', help='show a preview video with overlaid motion vectors')
+    parser.add_argument('-v', '--verbose', action='store_true', help='show detailled text output')
+    parser.add_argument('-d', '--dump', nargs='?', const=True,
+        help='dump frames, motion vectors, frame types, and timestamps to optionally specified output directory')
     args = parser.parse_args()
 
     if args.dump:
-        now = datetime.now().strftime("%Y-%m-%dT%H:%M:%S")
+        if isinstance(args.dump, str):
+            dumpdir = args.dump
+        else:
+            dumpdir = f"out-{datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}"
         for child in ["frames", "motion_vectors"]:
-            os.makedirs(os.path.join(f"out-{now}", child), exist_ok=True)
+            os.makedirs(os.path.join(dumpdir, child), exist_ok=True)
 
     cap = VideoCap()
 
@@ -83,14 +87,13 @@ def main(args=None):
 
         # store motion vectors, frames, etc. in output directory
         if args.dump:
-            cv2.imwrite(os.path.join(f"out-{now}", "frames", f"frame-{step}.jpg"), frame)
-            np.save(os.path.join(f"out-{now}", "motion_vectors", f"mvs-{step}.npy"), motion_vectors)
-            with open(os.path.join(f"out-{now}", "timestamps.txt"), "a") as f:
+            cv2.imwrite(os.path.join(dumpdir, "frames", f"frame-{step}.jpg"), frame)
+            np.save(os.path.join(dumpdir, "motion_vectors", f"mvs-{step}.npy"), motion_vectors)
+            with open(os.path.join(dumpdir, "timestamps.txt"), "a") as f:
                 f.write(str(timestamp)+"\n")
-            with open(os.path.join(f"out-{now}", "frame_types.txt"), "a") as f:
+            with open(os.path.join(dumpdir, "frame_types.txt"), "a") as f:
                 f.write(frame_type+"\n")
 
-
         step += 1
 
         if args.preview:

tests/end_to_end_tests.py

@@ -0,0 +1,65 @@
+import os
+import tempfile
+import unittest
+import subprocess
+
+import cv2
+import numpy as np
+
+
+PROJECT_ROOT = os.getenv("PROJECT_ROOT", "")
+
+
+class TestEndToEnd(unittest.TestCase):
+
+    def motions_vectors_valid(self, outdir, refdir):
+        equal = []
+        for i in range(336):
+            mvs = np.load(os.path.join(outdir, "motion_vectors", f"mvs-{i}.npy"))
+            mvs_ref = np.load(os.path.join(refdir, "motion_vectors", f"mvs-{i}.npy"))
+            equal.append(np.all(mvs == mvs_ref))
+        return all(equal)
+
+
+    def frame_types_valid(self, outdir, refdir):
+        with open(os.path.join(outdir, "frame_types.txt"), "r") as file:
+            frame_types = [line.strip() for line in file]
+        with open(os.path.join(refdir, "frame_types.txt"), "r") as file:
+            frame_types_ref = [line.strip() for line in file]
+        return frame_types == frame_types_ref
+
+
+    def frames_valid(self, outdir, refdir):
+        equal = []
+        for i in range(336):
+            frame = cv2.imread(os.path.join(outdir, "frames", f"frame-{i}.jpg"))
+            frame_ref = cv2.imread(os.path.join(refdir, "frames", f"frame-{i}.jpg"))
+            equal.append(np.all(frame == frame_ref))
+        return all(equal)
+
+
+    def test_end_to_end_h264(self):
+        with tempfile.TemporaryDirectory() as outdir:
+            print("Running extraction for H.264")
+            subprocess.run(f"extract_mvs {os.path.join(PROJECT_ROOT, 'vid_h264.mp4')} --dump {outdir}", shell=True, check=True)
+            refdir = os.path.join(PROJECT_ROOT, "tests/reference/h264")
+
+            self.assertTrue(self.motions_vectors_valid(outdir, refdir), msg="motion vectors are invalid")
+            self.assertTrue(self.frame_types_valid(outdir, refdir), msg="frame types are invalid")
+            self.assertTrue(self.frames_valid(outdir, refdir), msg="frames are invalid")
+
+
+    def test_end_to_end_mpeg4_part2(self):
+        with tempfile.TemporaryDirectory() as outdir:
+            print("Running extraction for MPEG-4 Part 2")
+            subprocess.run(f"extract_mvs {os.path.join(PROJECT_ROOT, 'vid_mpeg4_part2.mp4')} --dump {outdir}", shell=True, check=True)
+            refdir = os.path.join(PROJECT_ROOT, "tests/reference/mpeg4_part2")
+
+            self.assertTrue(self.motions_vectors_valid(outdir, refdir), msg="motion vectors are invalid")
+            self.assertTrue(self.frame_types_valid(outdir, refdir), msg="frame types are invalid")
+            self.assertTrue(self.frames_valid(outdir, refdir), msg="frames are invalid")
+
+
+if __name__ == '__main__':
+    unittest.main()
+