|
import numba as nb |
|
import numpy as np |
|
|
|
|
|
@nb.njit() |
|
def find_nearest_stft_bin(f0_, freqs): |
|
freqs = np.expand_dims(freqs, 0) |
|
f0_ = np.expand_dims(f0_, 1) |
|
return np.abs(freqs - f0_).argmin() |
|
|
|
|
|
@nb.njit() |
|
def get_med_curve(f0, step_size=20): |
|
v_begin = -1 |
|
v_end = -1 |
|
x_med_curve = [] |
|
y_med_curve = [] |
|
T = len(f0) |
|
|
|
for i in range(T): |
|
if f0[i] >= 50 and i < T - 1: |
|
if v_begin == -1: |
|
v_begin = i |
|
v_end = i |
|
else: |
|
if v_end != -1: |
|
if v_end - v_begin > 3: |
|
for j in range(v_begin, v_end + 1 - step_size, step_size): |
|
frag_med = np.median(f0[j:j + step_size]) |
|
x_med_curve.append(j) |
|
y_med_curve.append(frag_med) |
|
x_med_curve.append(v_end) |
|
y_med_curve.append(np.median(f0[v_end - step_size:v_end + 1])) |
|
v_end = v_begin = -1 |
|
x_med_curve = [0] + x_med_curve + [T] |
|
x_med_curve = np.array(x_med_curve) |
|
y_med_curve = [y_med_curve[0]] + y_med_curve + [y_med_curve[-1]] |
|
y_med_curve = np.array(y_med_curve) |
|
return x_med_curve, y_med_curve |
|
|
|
|
|
@nb.njit() |
|
def clean_short_v_frag(f0): |
|
v_begin = -1 |
|
T = len(f0) |
|
|
|
uv = np.zeros_like(f0).astype(np.bool_) |
|
for i in range(T): |
|
if f0[i] >= 1e-4 and i < T - 1: |
|
if v_begin == -1: |
|
v_begin = i |
|
else: |
|
if v_begin != -1: |
|
v_end = i if f0[i] >= 1e-4 else i - 1 |
|
if v_end - v_begin + 1 < 3: |
|
uv[v_begin:v_end + 1] = 1 |
|
v_begin = -1 |
|
return uv |
|
|
|
|
|
@nb.njit() |
|
def find_best_f0_using_har_energy(spec, pitches, freqs, hars, hars_mhalf, f0_min, f0_max): |
|
re = np.zeros_like(spec) |
|
T = len(spec) |
|
for i in range(T): |
|
spec_i = spec[i] |
|
for j, f0_j in enumerate(pitches[i]): |
|
if f0_j == 0 or f0_j < f0_min[i] or f0_j > f0_max[i]: |
|
continue |
|
mask = np.zeros((10000,)) |
|
mask_mhalf = np.zeros((10000,)) |
|
for mul in hars: |
|
b = find_nearest_stft_bin(np.array((f0_j * mul,)), freqs) |
|
for delta in range(-1, 2): |
|
mask[b + delta] = 1 |
|
for mul in hars_mhalf: |
|
b_mhalf = find_nearest_stft_bin(np.array((f0_j * (mul - 0.5),)), freqs) |
|
for delta in range(-1, 2): |
|
mask_mhalf[b_mhalf + delta] = 1 |
|
mask = mask[:len(spec_i)] |
|
mask_mhalf = mask_mhalf[:len(spec_i)] |
|
energy = (np.exp(spec_i) * mask).sum() / mask.sum() |
|
energy_mhalf = (np.exp(spec_i) * mask_mhalf).sum() / mask_mhalf.sum() |
|
re[i, j] = energy / energy_mhalf |
|
f0_2d_mask = 10000 * (re > 2) + 20000 * (re > 3) + np.expand_dims(np.arange(re.shape[1])[::-1], 0) |
|
f0_idx = np.zeros((T,), dtype=np.int_) |
|
for i in range(T): |
|
f0_idx[i] = f0_2d_mask[i].argmax() |
|
uv = re.sum(-1) == 0 |
|
|
|
f0 = np.zeros((T,)) |
|
for i in range(T): |
|
f0[i] = pitches[i, f0_idx[i]] |
|
f0 = f0 * (1 - uv) |
|
uv = clean_short_v_frag(f0) |
|
f0[uv] = 0 |
|
x_med_curve, y_med_curve = get_med_curve(f0) |
|
re = re * (re > 1.5) |
|
return re, f0, x_med_curve, y_med_curve |
|
|