add ffprobe check and proper environment variables

.github/workflows/test_unittests.yml

@@ -10,7 +10,7 @@ jobs:
  runs-on: ubuntu-latest
  strategy:
  matrix:
- ffmpeg-version: ["3.2.4", "3.4", "4.0.2", "4.1", "4.2", "4.3"]
+ ffmpeg-version: ["3.2.4", "3.4.2", "4.4.2", "5.1.2", "6.1.1", "7.0.0"]
 
  # Timeout: https://stackoverflow.com/a/59076067/4521646
  timeout-minutes: 10
@@ -29,8 +29,10 @@ jobs:
  - name: Install dependencies FFMPEG ${{ matrix.ffmpeg-version }}
  env:
  FFMPEG_INSTALL: ${{ matrix.pytorch-version }}
+ MP4BOX_PATH: /usr/bin/MP4Box
  run: |
  sudo apt-get -y install gpac
+ which MP4Box
  conda install -c conda-forge ffmpeg==${{ matrix.ffmpeg-version }}
  python -m pip install -e .['tests']
  python --version

examples/readwrite.py

@@ -9,117 +9,25 @@
 
 if __name__ == '__main__':
  parser = argparse.ArgumentParser()
- parser.add_argument(
- 'input',
- )
  args = parser.parse_args()
 
  # load stems
- stems, rate = stempeg.read_stems(args.input)
-
- # load stems,
- # resample to 96000 Hz,
- # use multiprocessing
- stems, rate = stempeg.read_stems(
- args.input,
- sample_rate=96000,
- multiprocess=True
- )
+ stems, rate = stempeg.read_stems(stempeg.example_stem_path())
 
- # --> stems now has `shape=(stem x samples x channels)``
-
- # save stems from tensor as multi-stream mp4
- stempeg.write_stems(
- "test.stem.m4a",
- stems,
- sample_rate=96000
- )
-
- # save stems as dict for convenience
  stems = {
  "mix": stems[0],
  "drums": stems[1],
  "bass": stems[2],
  "other": stems[3],
  "vocals": stems[4],
  }
- # keys will be automatically used
-
- # from dict as files
- stempeg.write_stems(
- "test.stem.m4a",
- data=stems,
- sample_rate=96000
- )
 
- # `write_stems` is a preset for the following settings
- # here the output signal is resampled to 44100 Hz and AAC codec is used
  stempeg.write_stems(
- "test.stem.m4a",
- stems,
- sample_rate=96000,
- writer=stempeg.StreamsWriter(
- codec="aac",
- output_sample_rate=44100,
- bitrate="256000",
- stem_names=['mix', 'drums', 'bass', 'other', 'vocals']
- )
- )
-
- # Native Instruments compatible stems
- stempeg.write_stems(
- "test_traktor.stem.m4a",
- stems,
- sample_rate=96000,
- writer=stempeg.NIStemsWriter(
- stems_metadata=[
- {"color": "#009E73", "name": "Drums"},
- {"color": "#D55E00", "name": "Bass"},
- {"color": "#CC79A7", "name": "Other"},
- {"color": "#56B4E9", "name": "Vocals"}
- ]
- )
- )
-
- # lets write as multistream opus (supports only 48000 khz)
- stempeg.write_stems(
- "test.stem.opus",
- stems,
- sample_rate=96000,
- writer=stempeg.StreamsWriter(
- output_sample_rate=48000,
- codec="opus"
- )
- )
-
- # writing to wav requires to convert streams to multichannel
- stempeg.write_stems(
- "test.wav",
- stems,
- sample_rate=96000,
- writer=stempeg.ChannelsWriter(
- output_sample_rate=48000
- )
- )
-
- # # stempeg also supports to load merged-multichannel streams using
- stems, rate = stempeg.read_stems(
- "test.wav",
- reader=stempeg.ChannelsReader(nb_channels=2)
- )
-
- # mp3 does not support multiple channels,
- # therefore we have to use `stempeg.FilesWriter`
- # outputs are named ["output/0.mp3", "output/1.mp3"]
- # for named files, provide a dict or use `stem_names`
- # also apply multiprocessing
- stempeg.write_stems(
- ("output", ".mp3"),
+ ("output", ".flac"),
  stems,
  sample_rate=rate,
  writer=stempeg.FilesWriter(
- multiprocess=True,
- output_sample_rate=48000,
+ output_sample_rate=44100,
  stem_names=["mix", "drums", "bass", "other", "vocals"]
  )
  )

stempeg/__init__.py

@@ -20,7 +20,7 @@
 from .write import write_audio
 from .write import FilesWriter, StreamsWriter, ChannelsWriter, NIStemsWriter
 
-from .cmds import check_available_aac_encoders
+from .cmds import check_available_aac_encoders, ffmpeg_exists, ffprobe_exists, mp4box_exists
 
 import re
 import os

stempeg/cmds.py

@@ -1,6 +1,7 @@
 import re
 import subprocess as sp
 import logging
+import os
 
 FFMPEG_PATH = None
 FFPROBE_PATH = None
@@ -20,30 +21,56 @@ def find_cmd(cmd):
  return None
 
 
-def ffmpeg_and_ffprobe_exists():
- global FFMPEG_PATH, FFPROBE_PATH
- if FFMPEG_PATH is None:
+def ffmpeg_exists():
+ global FFMPEG_PATH
+ # check environment variable
+ if "FFMPEG_PATH" in os.environ:
+ env_path = os.environ["FFMPEG_PATH"]
+ FFMPEG_PATH = find_cmd(env_path)
+ else:
  FFMPEG_PATH = find_cmd("ffmpeg")
 
- if FFPROBE_PATH is None:
+ return FFMPEG_PATH is not None
+
+
+def ffprobe_exists():
+ global FFPROBE_PATH
+ if "FFPROBE_PATH" in os.environ:
+ env_path = os.environ["FFPROBE_PATH"]
+ FFPROBE_PATH = find_cmd(env_path)
+ else:
  FFPROBE_PATH = find_cmd("ffprobe")
 
- return FFMPEG_PATH is not None and FFPROBE_PATH is not None
+ return FFPROBE_PATH is not None
 
 
 def mp4box_exists():
  global MP4BOX_PATH
- if MP4BOX_PATH is None:
+ print(MP4BOX_PATH)
+
+ if "MP4BOX_PATH" in os.environ:
+ env_path = os.environ["MP4BOX_PATH"]
+ MP4BOX_PATH = find_cmd(env_path)
+ else:
  MP4BOX_PATH = find_cmd("MP4Box")
+ print(MP4BOX_PATH)
 
  return MP4BOX_PATH is not None
 
 
-if not ffmpeg_and_ffprobe_exists():
+if not ffmpeg_exists():
+ raise RuntimeError(
+ "ffmpeg could not be found! "
+ "Please install it before using stempeg. "
+ "See: https://github.com/faroit/stempeg"
+ )
+
+
+if not ffprobe_exists():
  raise RuntimeError(
- 'ffmpeg or ffprobe could not be found! '
- 'Please install them before using stempeg. '
- 'See: https://github.com/faroit/stempeg'
+ "ffprobe could not be found! "
+ "Please install it before using stempeg. "
+ "See: https://github.com/faroit/stempeg"
  )
 
 

stempeg/write.py

@@ -20,7 +20,7 @@
 
 import stempeg
 
-from .cmds import FFMPEG_PATH, mp4box_exists, get_aac_codec, find_cmd
+from .cmds import FFMPEG_PATH, MP4BOX_PATH, mp4box_exists, get_aac_codec, find_cmd
 
 
 def _build_channel_map(nb_stems, nb_channels, stem_names=None):
@@ -497,7 +497,6 @@ def __init__(
  'Please install them before using NIStemsWriter().'
  'See: https://github.com/faroit/stempeg'
  )
- self.mp4boxcli = find_cmd("MP4Box")
  self.bitrate = bitrate
  self.default_metadata = default_metadata
  self.stems_metadata = stems_metadata
@@ -564,7 +563,7 @@ def __call__(
  if self.stems_metadata is not None:
  metadata['stems'] = self.stems_metadata
 
- callArgs = [self.mp4boxcli]
+ callArgs = [MP4BOX_PATH]
  callArgs.extend(["-add", str(Path(tempdir, '0.m4a#ID=Z')), path])
  for s in range(1, data.shape[0]):
  callArgs.extend(
@@ -758,7 +757,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
  )

tests/test_read.py

@@ -1,6 +1,7 @@
 import stempeg
 import numpy as np
 import pytest
+import os
 
 
 @pytest.fixture(params=[np.float16, np.float32, np.float64])
@@ -104,3 +105,15 @@ def test_info():
  fp = stempeg.example_stem_path()
  info = stempeg.Info(fp)
  S, rate = stempeg.read_stems(fp, info=info)
+
+
+def test_cmd():
+ assert stempeg.ffmpeg_exists()
+ assert stempeg.mp4box_exists()
+ assert stempeg.ffprobe_exists()
+ os.environ["FFMPEG_PATH"] = "/path_that_does_not_exist"
+ os.environ["FFPROBE_PATH"] = "/path_that_does_not_exist"
+ os.environ["MP4BOX_PATH"] = "/path_that_does_not_exist"
+ assert not stempeg.ffprobe_exists()
+ assert not stempeg.ffmpeg_exists()
+ assert not stempeg.mp4box_exists() 

tests/test_write.py

@@ -1,4 +1,5 @@
 from stempeg.write import ChannelsWriter
+from stempeg.cmds import MP4BOX_PATH 
 import stempeg
 import numpy as np
 import pytest
@@ -161,8 +162,6 @@ def ordered(obj):
 
 
 def test_nistems():
- mp4exc = stempeg.cmds.find_cmd("MP4Box")
-
  stems, rate = stempeg.read_stems(stempeg.example_stem_path())
  with tmp.NamedTemporaryFile(
  delete=False,
@@ -175,7 +174,7 @@ def test_nistems():
  sample_rate=rate,
  writer=stempeg.NIStemsWriter()
  )
- callArgs = [mp4exc]
+ callArgs = [MP4BOX_PATH]
  callArgs.extend(["-dump-udta", "0:stem", tempfile.name])
  sp.check_call(callArgs)