import os import json import argparse import numpy as np import matplotlib.pyplot as plt import torch import tqdm import librosa import librosa.display import soundfile as sf import pyloudnorm as pyln from dotmap import DotMap import gradio as gr from models import load_model_with_args from separate_func import ( conv_tasnet_separate, ) from utils import db2linear tqdm.monitor_interval = 0 def separate_track_with_model( args, model, device, track_audio, track_name, meter, augmented_gain ): with torch.no_grad(): if ( args.model_loss_params.architecture == "conv_tasnet_mask_on_output" or args.model_loss_params.architecture == "conv_tasnet" ): estimates = conv_tasnet_separate( args, model, device, track_audio, track_name, meter=meter, augmented_gain=augmented_gain, ) return estimates def main(input, mix_coefficient): parser = argparse.ArgumentParser(description="model test.py") parser.add_argument("--target", type=str, default="all") parser.add_argument("--weight_directory", type=str, default="weight") parser.add_argument("--output_directory", type=str, default="output") parser.add_argument("--use_gpu", type=bool, default=True) parser.add_argument("--save_name_as_target", type=bool, default=False) parser.add_argument( "--loudnorm_input_lufs", type=float, default=None, help="If you want to use loudnorm for input", ) parser.add_argument( "--save_output_loudnorm", type=float, default=-14.0, help="Save loudness normalized outputs or not. If you want to save, input target loudness", ) parser.add_argument( "--save_mixed_output", type=float, default=None, help="Save original+delimited-estimation mixed output with a ratio of default 0.5 (orginal) and 1 - 0.5 (estimation)", ) parser.add_argument( "--save_16k_mono", type=bool, default=False, help="Save 16k mono wav files for FAD evaluation.", ) parser.add_argument( "--save_histogram", type=bool, default=False, help="Save histogram of the output. Only valid when the task is 'delimit'", ) parser.add_argument( "--use_singletrackset", type=bool, default=False, help="Use SingleTrackSet if input data is too long.", ) args, _ = parser.parse_known_args() with open(f"{args.weight_directory}/{args.target}.json", "r") as f: args_dict = json.load(f) args_dict = DotMap(args_dict) for key, value in args_dict["args"].items(): if key in list(vars(args).keys()): pass else: setattr(args, key, value) args.test_output_dir = f"{args.output_directory}" os.makedirs(args.test_output_dir, exist_ok=True) device = torch.device( "cuda" if torch.cuda.is_available() and args.use_gpu else "cpu" ) ###################### Define Models ###################### our_model = load_model_with_args(args) our_model = our_model.to(device) target_model_path = f"{args.weight_directory}/{args.target}.pth" checkpoint = torch.load(target_model_path, map_location=device) our_model.load_state_dict(checkpoint) our_model.eval() meter = pyln.Meter(44100) sr, track_audio = input track_audio = track_audio.T track_name = "gradio_demo" orig_audio = track_audio.copy() if sr != 44100: raise ValueError("Sample rate should be 44100") augmented_gain = None if args.loudnorm_input_lufs: # If you want to use loud-normalized input track_lufs = meter.integrated_loudness(track_audio.T) augmented_gain = args.loudnorm_input_lufs - track_lufs track_audio = track_audio * db2linear(augmented_gain, eps=0.0) track_audio = ( torch.as_tensor(track_audio, dtype=torch.float32).unsqueeze(0).to(device) ) estimates = separate_track_with_model( args, our_model, device, track_audio, track_name, meter, augmented_gain ) if args.save_mixed_output: track_lufs = meter.integrated_loudness(orig_audio.T) augmented_gain = args.save_output_loudnorm - track_lufs orig_audio = orig_audio * db2linear(augmented_gain, eps=0.0) mixed_output = orig_audio * args.save_mixed_output + estimates * ( 1 - args.save_mixed_output ) sf.write( f"{args.test_output_dir}/{track_name}/{track_name}_mixed.wav", mixed_output.T, args.data_params.sample_rate, ) return ( (sr, estimates.T), (sr, orig_audio.T), (sr, orig_audio.T * mix_coefficient + estimates.T * (1 - mix_coefficient)), ) def parallel_mix(input, output, mix_coefficient): sr = 44100 return sr, input[1] * mix_coefficient + output[1] * (1 - mix_coefficient) def int16_to_float32(wav): wav = np.frombuffer(wav, dtype=np.int16) X = wav / 32768 return X def waveform_plot(input, output, prl_mix_ouptut, figsize_x=20, figsize_y=9): sr = 44100 fig, ax = plt.subplots( nrows=3, sharex=True, sharey=True, figsize=(figsize_x, figsize_y) ) librosa.display.waveshow(int16_to_float32(input[1]).T, sr=sr, ax=ax[0]) ax[0].set(title="Loudness Normalized Input") ax[0].label_outer() librosa.display.waveshow(int16_to_float32(output[1]).T, sr=sr, ax=ax[1]) ax[1].set(title="De-limiter Output") ax[1].label_outer() librosa.display.waveshow(int16_to_float32(prl_mix_ouptut[1]).T, sr=sr, ax=ax[2]) ax[2].set(title="Parallel Mix of the Input and its De-limiter Output") ax[2].label_outer() return fig with gr.Blocks() as demo: gr.HTML( """

Music De-limiter

A demo for "Music De-limiter via Sample-wise Gain Inversion" to appear in WASPAA 2023. You can first upload a music (.wav or .mp3) file and then press "De-limit" button to apply the De-limiter. Since we use a CPU instead of a GPU, it may require a few minute. Then, you can apply a Parallel Mix technique, which is a simple linear mixing technique of "loudness normalized input" and the "de-limiter output". You can modify the mixing coefficient by yourself. If the coefficient is 0.3 then the output will be the "loudness_normalized_input * 0.3 + de-limiter_output * 0.7"

""" ) with gr.Row().style(mobile_collapse=False, equal_height=True): with gr.Column(): with gr.Box(): input_audio = gr.Audio(source="upload", label="De-limiter Input") btn = gr.Button("De-limit") with gr.Column(): with gr.Box(): loud_norm_input = gr.Audio(label="Loudness Normalized Input (-14LUFS)") with gr.Box(): output_audio = gr.Audio(label="De-limiter Output") with gr.Box(): output_audio_parallel = gr.Audio( label="Parallel Mix of the Input and its De-limiter Output" ) slider = gr.Slider( minimum=0, maximum=1, step=0.1, value=0.5, label="Parallel Mix Coefficient", ) btn.click( main, inputs=[input_audio, slider], outputs=[output_audio, loud_norm_input, output_audio_parallel], ) slider.release( parallel_mix, inputs=[input_audio, output_audio, slider], outputs=output_audio_parallel, ) with gr.Row().style(mobile_collapse=False, equal_height=True): with gr.Column(): with gr.Box(): plot = gr.Plot(label="Plots") btn2 = gr.Button("Show Plots") slider_plot_x = gr.Slider( minimum=1, maximum=100, step=1, value=20, label="Plot X-axis size", ) slider_plot_y = gr.Slider( minimum=1, maximum=30, step=1, value=9, label="Plot Y-axis size", ) btn2.click( waveform_plot, inputs=[ loud_norm_input, output_audio, output_audio_parallel, slider_plot_x, slider_plot_y, ], outputs=plot, ) slider.release( waveform_plot, inputs=[ loud_norm_input, output_audio, output_audio_parallel, slider_plot_x, slider_plot_y, ], outputs=plot, ) if __name__ == "__main__": demo.launch(debug=True)