diff --git a/.dockerignore b/.dockerignore index 1fcfd93..f2b78e5 100644 --- a/.dockerignore +++ b/.dockerignore @@ -7,4 +7,5 @@ __pycache__ /config /config.yml pyproject.toml -poetry.lock \ No newline at end of file +poetry.lock +/tests \ No newline at end of file diff --git a/.github/workflows/_test.yml b/.github/workflows/_test.yml index 36d3401..b94a212 100644 --- a/.github/workflows/_test.yml +++ b/.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" diff --git a/.gitignore b/.gitignore index 0942676..d34f859 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,10 @@ -*.mp4 +/*.mp4 /config-local.yml /data/* !/data/README.md __pycache__ -.coverage \ No newline at end of file +.coverage +/tests/data/mp4s/* +!/tests/data/mp4s/test.mp4 +!/tests/data/spectograms/* +!/tests/data/spectograms/README.md \ No newline at end of file diff --git a/README.md b/README.md index 25234f2..9cc053e 100644 --- a/README.md +++ b/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) \ No newline at end of file +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 \ No newline at end of file diff --git a/base_util.py b/base_util.py index a22aa99..333891f 100644 --- a/base_util.py +++ b/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""" diff --git a/pyproject.toml b/pyproject.toml index 7716126..296f461 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -70,5 +70,6 @@ module = [ 'yaml', 'yacs.*', 'numpy', + 'pydub', ] ignore_missing_imports = true diff --git a/spectogram_util.py b/spectogram_util.py index f13f8da..127f5d2 100644 --- a/spectogram_util.py +++ b/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( diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/data/mp4s/test.mp4 b/tests/data/mp4s/test.mp4 new file mode 100644 index 0000000..0c43339 Binary files /dev/null and b/tests/data/mp4s/test.mp4 differ diff --git a/tests/data/spectograms/README.md b/tests/data/spectograms/README.md new file mode 100644 index 0000000..beb5511 --- /dev/null +++ b/tests/data/spectograms/README.md @@ -0,0 +1 @@ +used by the spectogram unit test to temporarily store spectograms generated in spectogram_util_test.py \ No newline at end of file diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/unit/spectogram_util_test.py b/tests/unit/spectogram_util_test.py new file mode 100644 index 0000000..ebe3e20 --- /dev/null +++ b/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) diff --git a/visxp_prep.py b/visxp_prep.py index 8c0f751..4c1f135 100644 --- a/visxp_prep.py +++ b/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)