juancopi81's picture
Add t5x and mt3 models
b100e1c
raw
history blame
15.4 kB
# Copyright 2022 The MT3 Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Transcription metrics."""
import collections
import copy
import functools
from typing import Any, Iterable, Mapping, Optional, Sequence
import mir_eval
from mt3 import event_codec
from mt3 import metrics_utils
from mt3 import note_sequences
from mt3 import spectrograms
from mt3 import summaries
from mt3 import vocabularies
import note_seq
import numpy as np
import seqio
def _program_aware_note_scores(
ref_ns: note_seq.NoteSequence,
est_ns: note_seq.NoteSequence,
granularity_type: str
) -> Mapping[str, float]:
"""Compute precision/recall/F1 for notes taking program into account.
For non-drum tracks, uses onsets and offsets. For drum tracks, uses onsets
only. Applies MIDI program map of specified granularity type.
Args:
ref_ns: Reference NoteSequence with ground truth labels.
est_ns: Estimated NoteSequence.
granularity_type: String key in vocabularies.PROGRAM_GRANULARITIES dict.
Returns:
A dictionary containing precision, recall, and F1 score.
"""
program_map_fn = vocabularies.PROGRAM_GRANULARITIES[
granularity_type].program_map_fn
ref_ns = copy.deepcopy(ref_ns)
for note in ref_ns.notes:
if not note.is_drum:
note.program = program_map_fn(note.program)
est_ns = copy.deepcopy(est_ns)
for note in est_ns.notes:
if not note.is_drum:
note.program = program_map_fn(note.program)
program_and_is_drum_tuples = (
set((note.program, note.is_drum) for note in ref_ns.notes) |
set((note.program, note.is_drum) for note in est_ns.notes)
)
drum_precision_sum = 0.0
drum_precision_count = 0
drum_recall_sum = 0.0
drum_recall_count = 0
nondrum_precision_sum = 0.0
nondrum_precision_count = 0
nondrum_recall_sum = 0.0
nondrum_recall_count = 0
for program, is_drum in program_and_is_drum_tuples:
est_track = note_sequences.extract_track(est_ns, program, is_drum)
ref_track = note_sequences.extract_track(ref_ns, program, is_drum)
est_intervals, est_pitches, unused_est_velocities = (
note_seq.sequences_lib.sequence_to_valued_intervals(est_track))
ref_intervals, ref_pitches, unused_ref_velocities = (
note_seq.sequences_lib.sequence_to_valued_intervals(ref_track))
args = {
'ref_intervals': ref_intervals, 'ref_pitches': ref_pitches,
'est_intervals': est_intervals, 'est_pitches': est_pitches
}
if is_drum:
args['offset_ratio'] = None
precision, recall, unused_f_measure, unused_avg_overlap_ratio = (
mir_eval.transcription.precision_recall_f1_overlap(**args))
if is_drum:
drum_precision_sum += precision * len(est_intervals)
drum_precision_count += len(est_intervals)
drum_recall_sum += recall * len(ref_intervals)
drum_recall_count += len(ref_intervals)
else:
nondrum_precision_sum += precision * len(est_intervals)
nondrum_precision_count += len(est_intervals)
nondrum_recall_sum += recall * len(ref_intervals)
nondrum_recall_count += len(ref_intervals)
precision_sum = drum_precision_sum + nondrum_precision_sum
precision_count = drum_precision_count + nondrum_precision_count
recall_sum = drum_recall_sum + nondrum_recall_sum
recall_count = drum_recall_count + nondrum_recall_count
precision = (precision_sum / precision_count) if precision_count else 0
recall = (recall_sum / recall_count) if recall_count else 0
f_measure = mir_eval.util.f_measure(precision, recall)
drum_precision = ((drum_precision_sum / drum_precision_count)
if drum_precision_count else 0)
drum_recall = ((drum_recall_sum / drum_recall_count)
if drum_recall_count else 0)
drum_f_measure = mir_eval.util.f_measure(drum_precision, drum_recall)
nondrum_precision = ((nondrum_precision_sum / nondrum_precision_count)
if nondrum_precision_count else 0)
nondrum_recall = ((nondrum_recall_sum / nondrum_recall_count)
if nondrum_recall_count else 0)
nondrum_f_measure = mir_eval.util.f_measure(nondrum_precision, nondrum_recall)
return {
f'Onset + offset + program precision ({granularity_type})': precision,
f'Onset + offset + program recall ({granularity_type})': recall,
f'Onset + offset + program F1 ({granularity_type})': f_measure,
f'Drum onset precision ({granularity_type})': drum_precision,
f'Drum onset recall ({granularity_type})': drum_recall,
f'Drum onset F1 ({granularity_type})': drum_f_measure,
f'Nondrum onset + offset + program precision ({granularity_type})':
nondrum_precision,
f'Nondrum onset + offset + program recall ({granularity_type})':
nondrum_recall,
f'Nondrum onset + offset + program F1 ({granularity_type})':
nondrum_f_measure
}
def _note_onset_tolerance_sweep(
ref_ns: note_seq.NoteSequence, est_ns: note_seq.NoteSequence,
tolerances: Iterable[float] = (0.01, 0.02, 0.05, 0.1, 0.2, 0.5)
) -> Mapping[str, float]:
"""Compute note precision/recall/F1 across a range of tolerances."""
est_intervals, est_pitches, unused_est_velocities = (
note_seq.sequences_lib.sequence_to_valued_intervals(est_ns))
ref_intervals, ref_pitches, unused_ref_velocities = (
note_seq.sequences_lib.sequence_to_valued_intervals(ref_ns))
scores = {}
for tol in tolerances:
precision, recall, f_measure, _ = (
mir_eval.transcription.precision_recall_f1_overlap(
ref_intervals=ref_intervals, ref_pitches=ref_pitches,
est_intervals=est_intervals, est_pitches=est_pitches,
onset_tolerance=tol, offset_min_tolerance=tol))
scores[f'Onset + offset precision ({tol})'] = precision
scores[f'Onset + offset recall ({tol})'] = recall
scores[f'Onset + offset F1 ({tol})'] = f_measure
return scores
def transcription_metrics(
targets: Sequence[Mapping[str, Any]],
predictions: Sequence[Mapping[str, Any]],
codec: event_codec.Codec,
spectrogram_config: spectrograms.SpectrogramConfig,
onsets_only: bool,
use_ties: bool,
track_specs: Optional[Sequence[note_sequences.TrackSpec]] = None,
num_summary_examples: int = 5,
frame_fps: float = 62.5,
frame_velocity_threshold: int = 30,
) -> Mapping[str, seqio.metrics.MetricValue]:
"""Compute mir_eval transcription metrics."""
if onsets_only and use_ties:
raise ValueError('Ties not compatible with onset-only transcription.')
if onsets_only:
encoding_spec = note_sequences.NoteOnsetEncodingSpec
elif not use_ties:
encoding_spec = note_sequences.NoteEncodingSpec
else:
encoding_spec = note_sequences.NoteEncodingWithTiesSpec
# The first target for each full example contains the NoteSequence; just
# organize by ID.
full_targets = {}
for target in targets:
if target['ref_ns']:
full_targets[target['unique_id']] = {'ref_ns': target['ref_ns']}
# Gather all predictions for the same ID and concatenate them in time order,
# to construct full-length predictions.
full_predictions = metrics_utils.combine_predictions_by_id(
predictions=predictions,
combine_predictions_fn=functools.partial(
metrics_utils.event_predictions_to_ns,
codec=codec,
encoding_spec=encoding_spec))
assert sorted(full_targets.keys()) == sorted(full_predictions.keys())
full_target_prediction_pairs = [
(full_targets[id], full_predictions[id])
for id in sorted(full_targets.keys())
]
scores = collections.defaultdict(list)
all_track_pianorolls = collections.defaultdict(list)
for target, prediction in full_target_prediction_pairs:
scores['Invalid events'].append(prediction['est_invalid_events'])
scores['Dropped events'].append(prediction['est_dropped_events'])
def remove_drums(ns):
ns_drumless = note_seq.NoteSequence()
ns_drumless.CopyFrom(ns)
del ns_drumless.notes[:]
ns_drumless.notes.extend([note for note in ns.notes if not note.is_drum])
return ns_drumless
est_ns_drumless = remove_drums(prediction['est_ns'])
ref_ns_drumless = remove_drums(target['ref_ns'])
# Whether or not there are separate tracks, compute metrics for the full
# NoteSequence minus drums.
est_tracks = [est_ns_drumless]
ref_tracks = [ref_ns_drumless]
use_track_offsets = [not onsets_only]
use_track_velocities = [not onsets_only]
track_instrument_names = ['']
if track_specs is not None:
# Compute transcription metrics separately for each track.
for spec in track_specs:
est_tracks.append(note_sequences.extract_track(
prediction['est_ns'], spec.program, spec.is_drum))
ref_tracks.append(note_sequences.extract_track(
target['ref_ns'], spec.program, spec.is_drum))
use_track_offsets.append(not onsets_only and not spec.is_drum)
use_track_velocities.append(not onsets_only)
track_instrument_names.append(spec.name)
for est_ns, ref_ns, use_offsets, use_velocities, instrument_name in zip(
est_tracks, ref_tracks, use_track_offsets, use_track_velocities,
track_instrument_names):
track_scores = {}
est_intervals, est_pitches, est_velocities = (
note_seq.sequences_lib.sequence_to_valued_intervals(est_ns))
ref_intervals, ref_pitches, ref_velocities = (
note_seq.sequences_lib.sequence_to_valued_intervals(ref_ns))
# Precision / recall / F1 using onsets (and pitches) only.
precision, recall, f_measure, avg_overlap_ratio = (
mir_eval.transcription.precision_recall_f1_overlap(
ref_intervals=ref_intervals,
ref_pitches=ref_pitches,
est_intervals=est_intervals,
est_pitches=est_pitches,
offset_ratio=None))
del avg_overlap_ratio
track_scores['Onset precision'] = precision
track_scores['Onset recall'] = recall
track_scores['Onset F1'] = f_measure
if use_offsets:
# Precision / recall / F1 using onsets and offsets.
precision, recall, f_measure, avg_overlap_ratio = (
mir_eval.transcription.precision_recall_f1_overlap(
ref_intervals=ref_intervals,
ref_pitches=ref_pitches,
est_intervals=est_intervals,
est_pitches=est_pitches))
del avg_overlap_ratio
track_scores['Onset + offset precision'] = precision
track_scores['Onset + offset recall'] = recall
track_scores['Onset + offset F1'] = f_measure
if use_velocities:
# Precision / recall / F1 using onsets and velocities (no offsets).
precision, recall, f_measure, avg_overlap_ratio = (
mir_eval.transcription_velocity.precision_recall_f1_overlap(
ref_intervals=ref_intervals,
ref_pitches=ref_pitches,
ref_velocities=ref_velocities,
est_intervals=est_intervals,
est_pitches=est_pitches,
est_velocities=est_velocities,
offset_ratio=None))
track_scores['Onset + velocity precision'] = precision
track_scores['Onset + velocity recall'] = recall
track_scores['Onset + velocity F1'] = f_measure
if use_offsets and use_velocities:
# Precision / recall / F1 using onsets, offsets, and velocities.
precision, recall, f_measure, avg_overlap_ratio = (
mir_eval.transcription_velocity.precision_recall_f1_overlap(
ref_intervals=ref_intervals,
ref_pitches=ref_pitches,
ref_velocities=ref_velocities,
est_intervals=est_intervals,
est_pitches=est_pitches,
est_velocities=est_velocities))
track_scores['Onset + offset + velocity precision'] = precision
track_scores['Onset + offset + velocity recall'] = recall
track_scores['Onset + offset + velocity F1'] = f_measure
# Calculate framewise metrics.
is_drum = all([n.is_drum for n in ref_ns.notes])
ref_pr = metrics_utils.get_prettymidi_pianoroll(
ref_ns, frame_fps, is_drum=is_drum)
est_pr = metrics_utils.get_prettymidi_pianoroll(
est_ns, frame_fps, is_drum=is_drum)
all_track_pianorolls[instrument_name].append((est_pr, ref_pr))
frame_precision, frame_recall, frame_f1 = metrics_utils.frame_metrics(
ref_pr, est_pr, velocity_threshold=frame_velocity_threshold)
track_scores['Frame Precision'] = frame_precision
track_scores['Frame Recall'] = frame_recall
track_scores['Frame F1'] = frame_f1
for metric_name, metric_value in track_scores.items():
if instrument_name:
scores[f'{instrument_name}/{metric_name}'].append(metric_value)
else:
scores[metric_name].append(metric_value)
# Add program-aware note metrics for all program granularities.
# Note that this interacts with the training program granularity; in
# particular granularities *higher* than the training granularity are likely
# to have poor metrics.
for granularity_type in vocabularies.PROGRAM_GRANULARITIES:
for name, score in _program_aware_note_scores(
target['ref_ns'], prediction['est_ns'],
granularity_type=granularity_type).items():
scores[name].append(score)
# Add (non-program-aware) note metrics across a range of onset/offset
# tolerances.
for name, score in _note_onset_tolerance_sweep(
ref_ns=ref_ns_drumless, est_ns=est_ns_drumless).items():
scores[name].append(score)
mean_scores = {k: np.mean(v) for k, v in scores.items()}
score_histograms = {'%s (hist)' % k: seqio.metrics.Histogram(np.array(v))
for k, v in scores.items()}
# Pick several examples to summarize.
targets_to_summarize, predictions_to_summarize = zip(
*full_target_prediction_pairs[:num_summary_examples])
# Compute audio summaries.
audio_summaries = summaries.audio_summaries(
targets=targets_to_summarize,
predictions=predictions_to_summarize,
spectrogram_config=spectrogram_config)
# Compute transcription summaries.
transcription_summaries = summaries.transcription_summaries(
targets=targets_to_summarize,
predictions=predictions_to_summarize,
spectrogram_config=spectrogram_config,
ns_feature_suffix='ns',
track_specs=track_specs)
pianorolls_to_summarize = {
k: v[:num_summary_examples] for k, v in all_track_pianorolls.items()
}
prettymidi_pianoroll_summaries = summaries.prettymidi_pianoroll(
pianorolls_to_summarize, fps=frame_fps)
return {
**mean_scores,
**score_histograms,
**audio_summaries,
**transcription_summaries,
**prettymidi_pianoroll_summaries,
}