Unit test for spectogram_util, mimicing Tengs example (#18)

* Unit test for spectogram_util, mimicing Tengs example

* kept input mp4; spectograms are deleted after each test run; reorganized quite some code

* pytest activated in GHA; ignoring tests dir in docker image

* sudo

* added --fix-missing for ffmpeg install

* debian

* disable ffmpeg

* apt-get update

---------

Co-authored-by: Sara Veldhoen <s.veldhoen@beeldengeluid.nl>
Co-authored-by: Jaap Blom <jblom@beeldengeluid.nl>

.dockerignore

@@ -7,4 +7,5 @@ __pycache__
 /config
 /config.yml
 pyproject.toml
-poetry.lock
+poetry.lock
+/tests

.github/workflows/_test.yml

@@ -20,12 +20,16 @@ jobs:
       run: "pip install --user poetry"
 
     - name: "Install dev environment"
-      run: " poetry install --no-interaction --no-ansi"
+      run: poetry install --no-interaction --no-ansi
 
-    # - name: "test"
-    #   run: |
-    #     cp config_test.yml config.yml
-    #     poetry run pytest
+    - name: install libgl1
+      run: sudo apt-get install -y libgl1
+
+    - name: install ffmpeg
+      run: sudo apt-get update && sudo apt-get install -y ffmpeg --fix-missing
+
+    - name: "test"
+      run: poetry run pytest
 
     - name: "flake8"
       run: "poetry run flake8"

.gitignore

@@ -1,6 +1,10 @@
-*.mp4
+/*.mp4
 /config-local.yml
 /data/*
 !/data/README.md
 __pycache__
-.coverage
+.coverage
+/tests/data/mp4s/*
+!/tests/data/mp4s/test.mp4
+!/tests/data/spectograms/*
+!/tests/data/spectograms/README.md

README.md

@@ -16,4 +16,13 @@ Installing `python-opencv` in a virtualenv, and thus not as a system package, co
 apt-get install libgl1
 ```
 
-(also see: https://stackoverflow.com/questions/64664094/i-cannot-use-opencv2-and-received-importerror-libgl-so-1-cannot-open-shared-obj)
+To make sure the unit-test work as well
+
+```
+apt-get install ffmpeg
+```
+
+Also see:
+https://stackoverflow.com/questions/64664094/i-cannot-use-opencv2-and-received-importerror-libgl-so-1-cannot-open-shared-obj
+
+https://docs.opencv.org/4.x/d2/de6/tutorial_py_setup_in_ubuntu.html

base_util.py

@@ -132,6 +132,11 @@ def __validate_parent_dirs(paths: list) -> None:
         raise (e)
 
 
+def get_source_id(input_file_path: str) -> str:
+    fn = os.path.basename(input_file_path)
+    return fn[0 : fn.rfind(".")] if "." in fn else fn
+
+
 # used for hecate
 def run_shell_command(cmd: str) -> bytes:
     """Run cmd and return stdout"""

pyproject.toml

@@ -70,5 +70,6 @@ module = [
   'yaml',
   'yacs.*',
   'numpy',
+  'pydub',
 ]
 ignore_missing_imports = true

spectogram_util.py

@@ -124,8 +124,8 @@ def raw_audio_to_spectograms(
             f"Spectogram is a np array with dimensions: {np.array(spectogram).shape}"
         )
         spec_path = os.path.join(location, f"{keyframe}.npz")
-        # out_dict = {"audio": spectogram}
-        np.savez(spec_path, audio=spectogram)  # out_dict as argument gives mypy error
+        out_dict = {"audio": spectogram}
+        np.savez(spec_path, out_dict)  # type: ignore
 
 
 def extract_audio_spectograms(

tests/data/spectograms/README.md

@@ -0,0 +1 @@
+used by the spectogram unit test to temporarily store spectograms generated in spectogram_util_test.py

tests/unit/spectogram_util_test.py

@@ -0,0 +1,142 @@
+import pytest
+import shutil
+from spectogram_util import extract_audio_spectograms, raw_audio_to_spectrogram
+import os
+from base_util import get_source_id
+import numpy as np
+
+MP4_INPUT_DIR = (
+    "./tests/data/mp4s"  # any file in this dir will be subjected to this test
+)
+SPECTOGRAM_OUTPUT_DIR = (
+    "./tests/data/spectograms"  # will be cleaned up after each test run
+)
+TMP_OUTPUT_PATH = "/tmp"  # should be available on most systems
+
+
+def generate_source_ids():
+    mp4_files = []
+    for root, dirs, files in os.walk(MP4_INPUT_DIR):
+        for f in files:
+            if f.find(".mp4") != -1:
+                mp4_files.append(get_source_id(f))
+    return mp4_files
+
+
+def to_output_dir(source_id: str) -> str:
+    return f"{SPECTOGRAM_OUTPUT_DIR}/{source_id}_example_output"
+
+
+def to_input_file(source_id: str) -> str:
+    return f"{MP4_INPUT_DIR}/{source_id}.mp4"
+
+
+def cleanup_output(source_id: str):
+    output_path = to_output_dir(source_id)
+    try:
+        shutil.rmtree(output_path)
+    except Exception:
+        return False
+    return True
+
+
+# Copied from https://github.com/beeldengeluid/dane-visual-feature-extraction-worker/blob/main/example.py
+def generate_example_output(media_file: str, output_path: str):
+    import wave
+    import numpy as np
+    from pydub import AudioSegment
+
+    # Convert MP4 to WAV
+    audio = AudioSegment.from_file(media_file)
+    audio.set_frame_rate(48000)
+    wav_fn = os.path.join(output_path, "output.wav")
+    audio.export(wav_fn, format="wav")
+
+    # Read WAV file
+    wav_file = wave.open(wav_fn)
+    n_channels, sampwidth, framerate, n_frames = wav_file.getparams()[:4]
+    data = wav_file.readframes(n_frames)
+    raw_audio = np.frombuffer(data, dtype=np.int16)
+    raw_audio = raw_audio.reshape((n_channels, n_frames), order="F")
+    raw_audio = raw_audio.astype(np.float32) / 32768.0  # type: ignore
+
+    # Segment audio into 1 second chunks
+    n_samples = raw_audio.shape[1]
+    n_samples_per_second = 48000
+    n_samples_per_chunk = n_samples_per_second
+    n_chunks = int(n_samples / n_samples_per_chunk)
+    chunks = []
+    for i in range(n_chunks):
+        chunks.append(
+            raw_audio[:, i * n_samples_per_chunk : (i + 1) * n_samples_per_chunk]
+        )
+
+    # Compute spectrogram for each chunk
+    spectrograms = []
+    for chunk in chunks:
+        spectrograms.append(raw_audio_to_spectrogram(chunk))
+
+    # Save spectrogram to file
+    for i, spectrogram in enumerate(spectrograms):
+        spec_path = os.path.join(output_path, f"{i}.npz")
+        out_dict = {"audio": spectrogram}
+        np.savez(spec_path, out_dict)  # type: ignore
+
+
+def assert_example_output(output_path: str, n_files: int):
+    if not os.path.exists(output_path):
+        os.makedirs(output_path)
+    for i in range(n_files):
+        if not os.path.isfile(os.path.join(output_path, f"{i}.npz")):
+            return False
+    return True
+
+
+@pytest.mark.parametrize(
+    "source_id, keyframe_timestamps, tmp_location",
+    [
+        (
+            source_id,
+            [
+                500,
+                1500,
+                2500,
+                3500,
+                4500,
+                5500,
+                6500,
+                7500,
+                8500,
+            ],  # for now the same for each mp4
+            TMP_OUTPUT_PATH,
+        )
+        for source_id in generate_source_ids()
+    ],
+)
+def test_extract_audio_spectograms(
+    source_id: str, keyframe_timestamps: list, tmp_location: str
+):
+    media_file = to_input_file(source_id)
+    output_path = to_output_dir(source_id)
+
+    if not assert_example_output(output_path, len(keyframe_timestamps)):
+        generate_example_output(media_file, output_path)
+
+    extract_audio_spectograms(
+        media_file=media_file,
+        keyframe_timestamps=keyframe_timestamps,
+        location=tmp_location,
+        tmp_location=tmp_location,
+    )
+    for i, timestamp in enumerate(keyframe_timestamps):
+        # Load example spectogram (following https://github.com/beeldengeluid/dane-visual-feature-extraction-worker/blob/main/example.py)
+        example_path = os.path.join(output_path, f"{i}.npz")
+        example_data = np.load(example_path, allow_pickle=True)
+        example_spectogram = example_data["arr_0"].item()["audio"]
+
+        real_path = os.path.join(tmp_location, f"{timestamp}.npz")
+        real_data = np.load(real_path, allow_pickle=True)
+        real_spectogram = real_data["arr_0"].item()["audio"]
+
+        assert np.equal(real_spectogram, example_spectogram).all()
+    assert cleanup_output(source_id)

visxp_prep.py

@@ -5,6 +5,7 @@
 import spectogram_util
 import cv2  # type: ignore
 import os
+from base_util import get_source_id
 from dane.config import cfg
 from dataclasses import dataclass
 
@@ -24,7 +25,7 @@ def generate_input_for_feature_extraction(
 
     output_dirs = {}
     for kind in ["keyframes", "metadata", "spectograms", "tmp"]:
-        output_dir = os.path.join("/data", _get_source_id(input_file_path), kind)
+        output_dir = os.path.join("/data", get_source_id(input_file_path), kind)
         if not os.path.isdir(output_dir):
             os.makedirs(output_dir)
         output_dirs[kind] = output_dir
@@ -106,11 +107,6 @@ def generate_input_for_feature_extraction(
     return VisXPFeatureExtractionInput(500, "Not implemented yet!", -1)
 
 
-def _get_source_id(input_file_path: str) -> str:
-    fn = os.path.basename(input_file_path)
-    return fn[0 : fn.rfind(".")] if "." in fn else fn
-
-
 def _frame_index_to_timecode(frame_index: int, fps: float, out_format="ms"):
     if out_format == "ms":
         return round(frame_index / fps * 1000)