def main(conf): compute_metrics = update_compute_metrics(conf["compute_wer"], COMPUTE_METRICS) anno_df = pd.read_csv( Path(conf["test_dir"]).parent.parent.parent / "test_annotations.csv") wer_tracker = (MockWERTracker() if not conf["compute_wer"] else WERTracker( ASR_MODEL_PATH, anno_df)) model_path = os.path.join(conf["exp_dir"], "best_model.pth") model = DPRNNTasNet.from_pretrained(model_path) # Handle device placement if conf["use_gpu"]: model.cuda() model_device = next(model.parameters()).device test_set = LibriMix( csv_dir=conf["test_dir"], task=conf["task"], sample_rate=conf["sample_rate"], n_src=conf["train_conf"]["data"]["n_src"], segment=None, return_id=True, ) # Uses all segment length # Used to reorder sources only loss_func = PITLossWrapper(pairwise_neg_sisdr, pit_from="pw_mtx") # Randomly choose the indexes of sentences to save. eval_save_dir = os.path.join(conf["exp_dir"], conf["out_dir"]) ex_save_dir = os.path.join(eval_save_dir, "examples/") if conf["n_save_ex"] == -1: conf["n_save_ex"] = len(test_set) save_idx = random.sample(range(len(test_set)), conf["n_save_ex"]) series_list = [] torch.no_grad().__enter__() for idx in tqdm(range(len(test_set))): # Forward the network on the mixture. mix, sources, ids = test_set[idx] mix, sources = tensors_to_device([mix, sources], device=model_device) est_sources = model(mix.unsqueeze(0)) loss, reordered_sources = loss_func(est_sources, sources[None], return_est=True) mix_np = mix.cpu().data.numpy() sources_np = sources.cpu().data.numpy() est_sources_np = reordered_sources.squeeze(0).cpu().data.numpy() # For each utterance, we get a dictionary with the mixture path, # the input and output metrics utt_metrics = get_metrics( mix_np, sources_np, est_sources_np, sample_rate=conf["sample_rate"], metrics_list=COMPUTE_METRICS, ) utt_metrics["mix_path"] = test_set.mixture_path est_sources_np_normalized = normalize_estimates(est_sources_np, mix_np) utt_metrics.update(**wer_tracker( mix=mix_np, clean=sources_np, estimate=est_sources_np_normalized, wav_id=ids, sample_rate=conf["sample_rate"], )) series_list.append(pd.Series(utt_metrics)) # Save some examples in a folder. Wav files and metrics as text. if idx in save_idx: local_save_dir = os.path.join(ex_save_dir, "ex_{}/".format(idx)) os.makedirs(local_save_dir, exist_ok=True) sf.write(local_save_dir + "mixture.wav", mix_np, conf["sample_rate"]) # Loop over the sources and estimates for src_idx, src in enumerate(sources_np): sf.write(local_save_dir + "s{}.wav".format(src_idx), src, conf["sample_rate"]) for src_idx, est_src in enumerate(est_sources_np_normalized): sf.write( local_save_dir + "s{}_estimate.wav".format(src_idx), est_src, conf["sample_rate"], ) # Write local metrics to the example folder. with open(local_save_dir + "metrics.json", "w") as f: json.dump(utt_metrics, f, indent=0) # Save all metrics to the experiment folder. all_metrics_df = pd.DataFrame(series_list) all_metrics_df.to_csv(os.path.join(eval_save_dir, "all_metrics.csv")) # Print and save summary metrics final_results = {} for metric_name in compute_metrics: input_metric_name = "input_" + metric_name ldf = all_metrics_df[metric_name] - all_metrics_df[input_metric_name] final_results[metric_name] = all_metrics_df[metric_name].mean() final_results[metric_name + "_imp"] = ldf.mean() print("Overall metrics :") pprint(final_results) if conf["compute_wer"]: print("\nWER report") wer_card = wer_tracker.final_report_as_markdown() print(wer_card) # Save the report with open(os.path.join(eval_save_dir, "final_wer.md"), "w") as f: f.write(wer_card) with open(os.path.join(eval_save_dir, "final_metrics.json"), "w") as f: json.dump(final_results, f, indent=0) model_dict = torch.load(model_path, map_location="cpu") os.makedirs(os.path.join(conf["exp_dir"], "publish_dir"), exist_ok=True) publishable = save_publishable( os.path.join(conf["exp_dir"], "publish_dir"), model_dict, metrics=final_results, train_conf=train_conf, )
def main(conf): compute_metrics = update_compute_metrics(conf["compute_wer"], COMPUTE_METRICS) anno_df = pd.read_csv(Path(conf["test_dir"]).parent.parent.parent / "test_annotations.csv") wer_tracker = ( MockWERTracker() if not conf["compute_wer"] else WERTracker(ASR_MODEL_PATH, anno_df)