diff --git a/.github/workflows/test_unittests.yml b/.github/workflows/test_unittests.yml index 2104f87..0d84d3c 100644 --- a/.github/workflows/test_unittests.yml +++ b/.github/workflows/test_unittests.yml @@ -30,6 +30,7 @@ jobs: env: FFMPEG_INSTALL: ${{ matrix.pytorch-version }} run: | + sudo apt-get update sudo apt-get -y install gpac conda install -c conda-forge ffmpeg==${{ matrix.ffmpeg-version }} python -m pip install -e .['tests'] @@ -40,4 +41,4 @@ jobs: run: conda list - name: Run tests run: | - py.test tests -v \ No newline at end of file + py.test tests -v diff --git a/.gitignore b/.gitignore index fe1d1d7..9e213ed 100644 --- a/.gitignore +++ b/.gitignore @@ -7,6 +7,8 @@ envffmpeg/ .vscode/ .circleci/ +# Pycharm +.idea/ #####=== OSX ===##### .DS_Store diff --git a/stempeg/read.py b/stempeg/read.py index bbc90b0..fea01c4 100644 --- a/stempeg/read.py +++ b/stempeg/read.py @@ -4,6 +4,11 @@ """ +import codecs +import json +import os +import subprocess + from stempeg.write import FilesWriter import numpy as np import warnings @@ -13,6 +18,8 @@ import atexit from functools import partial import datetime as dt +from .cmds import mp4box_exists, find_cmd + class Reader(object): """Base class for reader @@ -287,7 +294,7 @@ def read_stems( ] stem_durations = np.array([t.shape[0] for t in stems]) if not (stem_durations == stem_durations[0]).all(): - warnings.warning("Stems differ in length and were shortend") + warnings.warn("Stems differ in length and were shortend") min_length = np.min(stem_durations) stems = [t[:min_length, :] for t in stems] @@ -321,6 +328,11 @@ def __init__(self, filename): stream for stream in self.info['streams'] if stream['codec_type'] == 'audio' ] + # Try to get the stem titles using MP4Box if possible, otherwise fall back to numbered stems + stem_titles = _read_mp4box_stem_titles(filename) + if stem_titles is not None: + for i, stream in enumerate(self.audio_streams): + stream['tags']['handler_name'] = stem_titles[i] @property def nb_audio_streams(self): @@ -385,3 +397,38 @@ def channels(self, idx): def __repr__(self): """Print stream information""" return pprint.pformat(self.audio_streams) + + +def _read_mp4box_stem_titles(filename): + """Reads a mp4 stem titles file using MP4Box + Mainly taken from https://github.com/axeldelafosse/stemgen/blob/master/ni-stem/_internal.py + """ + stem_titles = None + if mp4box_exists(): + + mp4box = find_cmd("MP4Box") + + try: + callArgs = [mp4box] + callArgs.extend(["-dump-udta", "0:stem", filename, '-quiet']) + subprocess.check_call(callArgs) + + except subprocess.CalledProcessError as e: + return None + + try: + root, ext = os.path.splitext(filename) + udtaFile = root + "_stem.udta" + fileObj = codecs.open(udtaFile, encoding="utf-8") + # fileObj.seek(8) # Not sure why in the original code this is needed? + _metadata = json.load(fileObj) + os.remove(udtaFile) + + # add the mixture stem first since its index is 0 as per rest of the project + stem_titles = ['Mixture'] + [d['name'] for d in _metadata['stems']] + + except FileNotFoundError as e: + return None + + return stem_titles + diff --git a/stempeg/write.py b/stempeg/write.py index d0b7880..198acd6 100644 --- a/stempeg/write.py +++ b/stempeg/write.py @@ -758,7 +758,7 @@ def write_stems( """ # check if ffmpeg installed if int(stempeg.ffmpeg_version()[0]) < 3: - warnings.warning( + warnings.warn( "Writing stems with FFMPEG version < 3 is unsupported", UserWarning ) diff --git a/tests/test_read.py b/tests/test_read.py index 5979337..51ec996 100644 --- a/tests/test_read.py +++ b/tests/test_read.py @@ -104,3 +104,9 @@ def test_info(): fp = stempeg.example_stem_path() info = stempeg.Info(fp) S, rate = stempeg.read_stems(fp, info=info) + + +def test_info_stem_titles(): + fp = stempeg.example_stem_path() + info = stempeg.Info(fp) + assert info.title_streams == ['Mixture', 'Drums', 'Bass', 'Other', 'Vox']