beatalignment / scripts /dataset-stats.py
william590y's picture
Upload folder using huggingface_hub
151b875 verified
Raw
History Blame Contribute Delete
2.65 kB
from argparse import ArgumentParser
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from anticipation.vocab import MIDI_TIME_OFFSET, MIDI_START_OFFSET, TIME_RESOLUTION, SEPARATOR
from anticipation.ops import max_time
plt.rcParams['font.family'] = 'serif'
plt.rcParams['font.serif'] = ['Computer Modern']
plt.rcParams['font.size'] = 16
def loghist(filename, data, title, xlabel):
sns.set_style('whitegrid')
plt.clf()
plt.figure(figsize=(10,4))
#plt.title(title)
plt.xscale('log')
plt.xlabel(xlabel)
plt.ylabel('Density')
plt.grid(True, which='both', linestyle='-', linewidth=0.5)
density = sns.kdeplot(data, bw_adjust=1.0)
plt.tight_layout()
fig = density.get_figure()
fig.savefig(filename, dpi=300)
if __name__ == '__main__':
parser = ArgumentParser(description='calculate statistics of a tokenized MIDI dataset')
parser.add_argument('-f', '--filename',
help='file containing a tokenized MIDI dataset')
parser.add_argument('-i', '--interarrival',
action='store_true',
help='request interarrival-time enocoding (default to arrival-time encoding)')
args = parser.parse_args()
print(f'Calculating statistics for {args.filename}')
time_lengths = []
token_counts = []
with open(args.filename, 'r') as f:
for i,line in tqdm(list(enumerate(f))):
if i % 10 != 0: continue
tokens = [int(token) for token in line.split()]
if args.interarrival:
time_lengths.append(sum(t-MIDI_TIME_OFFSET for t in tokens if t < MIDI_START_OFFSET))
token_counts.append(len(tokens))
else:
if SEPARATOR in tokens:
continue # counts are weird; just skip these
time_lengths.append(max_time(tokens[1:], seconds=False))
token_counts.append(len(tokens[1:]))
tokens_per_second = [TIME_RESOLUTION*tokens/float(time) for (tokens, time) in zip(token_counts, time_lengths)]
print('Total tokens:', sum(token_counts))
print(f'Total time: {float(sum(time_lengths))/(3600*TIME_RESOLUTION)} hours')
print('Mean tokens-per-second:', TIME_RESOLUTION*sum(token_counts)/float(sum(time_lengths)))
print('Std tokens-per-second:', np.std(tokens_per_second))
print(np.mean(tokens_per_second))
loghist('output/tokens_per_second.png',
tokens_per_second,
'Distribution of Tokens per Second',
'Tokens per Second (log10 scale)')