def main(args):
# MODEL
num_features = [args.features*i for i in range(1, args.levels+1)] if args.feature_growth == "add" else \
[args.features*2**i for i in range(0, args.levels)]
target_outputs = int(args.output_size * args.sr)
model = Waveunet(args.channels,
num_features,
args.channels,
args.instruments,
kernel_size=args.kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res,
separate=args.separate)
if args.cuda:
model = utils.DataParallel(model)
print("move model to gpu")
model.cuda()
print("Loading model from checkpoint " + str(args.load_model))
state = utils.load_model(model, None, args.load_model)
print('Step', state['step'])
preds = predict_song(args, args.input, model)
output_folder = os.path.dirname(
args.input) if args.output is None else args.output
for inst in preds.keys():
utils.write_wav(
os.path.join(output_folder,
os.path.basename(args.input) + "_" + inst + ".wav"),
preds[inst], args.sr)
def predict(self, input):
"""Separate tracks from input mixture audio"""
out_path = Path(tempfile.mkdtemp())
zip_path = Path(tempfile.mkdtemp()) / "output.zip"
preds = predict_song(self.args, input, self.model)
out_names = []
for inst in preds.keys():
temp_n = os.path.join(
str(out_path),
os.path.basename(str(input)) + "_" + inst + ".wav")
data.utils.write_wav(temp_n, preds[inst], self.args.sr)
out_names.append(temp_n)
with zipfile.ZipFile(str(zip_path), "w") as zf:
for i in out_names:
zf.write(str(i))
return zip_path
[args.features*2**i for i in range(0, args.levels)]
target_outputs = int(args.output_size * args.sr)
model = Waveunet(args.channels,
num_features,
args.channels,
INSTRUMENTS,
kernel_size=args.kernel_size,
target_output_size=target_outputs,
depth=args.depth,
strides=args.strides,
conv_type=args.conv_type,
res=args.res,
separate=args.separate)
if args.cuda:
model = utils.DataParallel(model)
print("move model to gpu")
model.cuda()
print("Loading model from checkpoint " + str(args.load_model))
state = utils.load_model(model, None, args.load_model)
preds = predict_song(args, args.input, model)
output_folder = os.path.dirname(
args.input) if args.output is None else args.output
for inst in preds.keys():
utils.write_wav(
os.path.join(output_folder,
os.path.basename(args.input) + "_" + inst + ".wav"),
preds[inst], args.sr)