demucsim.py

#!/bin/env python3
# pylint: disable=line-too-long
"""
Demucsim - די מקסים

A wrapper around demucs[1], an music separation tool; it can separate drums,
vocals, and bass (these are all called stems) from a given audio track.

The only problem with Demucs is its insufficient CLI:
1. Most of the time, I want only a small section of the track, so I don't want
to waste extra time processing portions I won't use!
2. It shoves the output into a nested directory, and also includes a redundant
file of everything that _wasn't_ extracted. Why

So this is a wrapper that better suits my workflow: I tell it what portion I do
want, and what stem I do want, and it places the output in the directory I'm
currently in as a single WAV file. Slick.

#### INSTALLATION ####
You'll need `demucs`[1] and `ffmpeg` available on PATH. Then, to install this
script:
```
wget -O ~/.local/bin/demucsim https://gist.githubusercontent.com/it-is-wednesday/aa413368e6bbab650568000e2b88bb09/raw/demucsim.py
chmod +x ~/.local/bin/demucsim
```

[1]: https://github.com/facebookresearch/demucs
"""
import argparse
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from shutil import copy2, which
from subprocess import run
from tempfile import TemporaryDirectory
from typing import Literal, Tuple

# time.strptime default to year 1900, but we want epoch :)
SECS_BETWEEN_1900_AND_EPOCH = 2208997240


@dataclass
class CliArgs:
    audio_file: Path
    time_range: str
    stem: Literal["bass", "drums", "vocals", "other"]
    outdir: Path
    invert: bool


def main():
    """Entry point"""
    ensure_exec_available("ffmpeg")
    ensure_exec_available("demucs")

    args = get_args()
    inpt = Path(args.audio_file).absolute()

    startpoint, duration = parse_time_range(args.time_range)
    range_args = ["-ss", str(startpoint), "-t", str(duration)]

    # Demucs's output
    # with the --two-stems parameter, demucs produces two files: {stem}.wav and
    # no_{stem}.wav
    demucs_out = f"no_{args.stem}" if args.invert else args.stem

    # our program's output
    output_filename = f"{args.audio_file.stem}_{args.time_range}_{demucs_out}.wav"

    with TemporaryDirectory() as tmpdir:
        tmp_path = f"{tmpdir}/out.wav"
        run(["ffmpeg", "-i", inpt, *range_args, tmp_path], check=True)
        run(["demucs", "-o", f"{tmpdir}/out", "--two-stems", args.stem, tmp_path], check=True)
        output = next(Path(tmpdir).glob(f"out/*/*/{demucs_out}.wav"))
        copy2(output, args.outdir / output_filename)

    print("\n 𓆏𓆏𓆏𓆏𓆏𓆏𓆏𓆏𓆏𓆏𓆏𓆏\n")
    print(f"Demucsim is done! Output is at {args.outdir / output_filename}")


def parse_timestamp(stamp: str) -> int:
    """
    >>> parse_timestamp("00:00")
    0
    >>> parse_timestamp("10:00")
    600
    >>> parse_timestamp("01:00:00")
    3600
    >>> parse_timestamp("1:0:0")
    3600
    """
    fmt = "%M:%S" if stamp.count(":") == 1 else "%H:%M:%S"
    return round(time.mktime(time.strptime(stamp, fmt)) + SECS_BETWEEN_1900_AND_EPOCH)


def parse_time_range(range_str: str) -> Tuple[int, int]:
    """Translates the range into ffmpeg language, returns a tuple of `-ss` and `-t`.

    In ffmpeg, the selected time is expressed in terms of startpoint
    (-ss flag), and duration (-t flag).

    The startpoint is the num of seconds between the beginning of the track and
    the beginning of the range, and the duration is the num of seconds between
    the two edges of the range.

    >>> parse_time_range("00:43-01:53")
    (43, 70)
    """
    ts_beginning, ts_end = range_str.split("-")
    startpoint = parse_timestamp(ts_beginning)
    duration = parse_timestamp(ts_end) - startpoint
    return startpoint, duration


def get_args() -> CliArgs:
    parser = argparse.ArgumentParser()
    parser.add_argument("audio_file", type=Path)
    parser.add_argument("time_range", help="e.g. 00:43-01:52")
    parser.add_argument("stem", choices=["bass", "drums", "vocals", "other"])
    parser.add_argument(
        "-o", "--outdir", help="Directory to output into", type=Path, default=Path(".")
    )
    parser.add_argument("-v", "--invert", help="Remove chosen stem", action="store_true")
    return CliArgs(**vars(parser.parse_args()))


def ensure_exec_available(executable: str):
    """Exits if exec isn't in PATH"""
    if which(executable) is None:
        print(f"I'll need you to install {executable}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()