rmc/batch_encode.py at main · buk0vec/rmc · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
"""
batch_encode.py

Encodes a whole lotta stuff for listening tests
"""

import argparse
from pathlib import Path

from features import RMCFeatures
from xrmc import Decode, Encode, detect_tempo

KBPS = [64, 96]
CONFIGS = {
    "base": RMCFeatures(),
    "pred": RMCFeatures(PREDICTION=True),
    "bs": RMCFeatures(BLOCK_SWITCHING=True),
    "pred_bs": RMCFeatures(PREDICTION=True, BLOCK_SWITCHING=True),
}


# Put known tempos of material here (filename only, full path not needed)
# e.g. "ringnoord.wav": 164
KNOWN_TEMPOS = {}

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        prog="batch_encode",
    )
    parser.add_argument("input_folder", type=str)
    parser.add_argument("-o", "--output", type=str, default="batch_encode")
    args = parser.parse_args()
    input_folder = Path(args.input_folder)
    output_folder = Path(args.output)
    if not input_folder.is_dir():
        raise SystemExit(f"Error: {args.input_folder} is not a valid directory")
    codec_folder = output_folder / "encoded"
    codec_folder.mkdir(parents=True, exist_ok=True)
    audio_folder = output_folder / "decoded"
    audio_folder.mkdir(exist_ok=True)
    total_audio_files = len(
        [x for x in input_folder.iterdir() if x.is_file() and x.suffix == ".wav"]
    )
    total_encoded_files = len(KBPS) * len(CONFIGS) * total_audio_files
    print(f"Encoding/decoding {total_encoded_files} total files...")
    file_count = 0
    for file in input_folder.iterdir():
        if file.is_dir() or file.suffix != ".wav":
            continue
        file_count += 1
        print(f"=== {file_count}/{total_audio_files} {file.name} ===")
        tempo = KNOWN_TEMPOS.get(file.name) or detect_tempo(str(file))
        current_codec_folder: Path = codec_folder / file.stem
        current_codec_folder.mkdir(exist_ok=True)
        current_audio_folder: Path = audio_folder / file.stem
        current_audio_folder.mkdir(exist_ok=True)
        for bitrate in KBPS:
            print(f"-> Starting {bitrate} KBPS...")
            for config in CONFIGS:
                stem = f"{file.stem}_{config}_{bitrate}"
                coded_path = current_codec_folder / f"{stem}.rmc"
                audio_path = current_audio_folder / f"{stem}.wav"

                Encode(
                    inFilename=str(file),
                    features=CONFIGS[config],
                    codedFilename=str(coded_path),
                    kbps=bitrate,
                    tempo=tempo,
                    verbose=True,
                )

                Decode(
                    inFilename=str(coded_path),
                    features=CONFIGS[config],
                    outFilename=str(audio_path),
                    verbose=True,
                )

    print("Done.")