mirror of
https://github.com/macaodha/batdetect2.git
synced 2025-06-29 14:41:58 +02:00
292 lines
7.1 KiB
Python
292 lines
7.1 KiB
Python
"""Test suite for feature extraction functions."""
|
|
|
|
import logging
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import pytest
|
|
|
|
import batdetect2.detector.compute_features as feats
|
|
from batdetect2 import api, types
|
|
from batdetect2.utils import audio_utils as au
|
|
|
|
numba_logger = logging.getLogger("numba")
|
|
numba_logger.setLevel(logging.WARNING)
|
|
|
|
|
|
def index_to_freq(
|
|
index: int,
|
|
spec_height: int,
|
|
min_freq: int,
|
|
max_freq: int,
|
|
) -> float:
|
|
"""Convert spectrogram index to frequency in Hz."""
|
|
index = spec_height - index
|
|
return round(
|
|
(index / float(spec_height)) * (max_freq - min_freq) + min_freq, 2
|
|
)
|
|
|
|
|
|
def index_to_time(
|
|
index: int,
|
|
spec_width: int,
|
|
spec_duration: float,
|
|
) -> float:
|
|
"""Convert spectrogram index to time in seconds."""
|
|
return round((index / float(spec_width)) * spec_duration, 2)
|
|
|
|
|
|
def test_get_feats_function_with_empty_spectrogram():
|
|
"""Test get_feats function with empty spectrogram.
|
|
|
|
This tests that the overall flow of the function works, even if the
|
|
spectrogram is empty.
|
|
"""
|
|
spec_duration = 3
|
|
spec_width = 100
|
|
spec_height = 100
|
|
min_freq = 10_000
|
|
max_freq = 120_000
|
|
spectrogram = np.zeros((spec_height, spec_width))
|
|
|
|
x_pos = 20
|
|
y_pos = 80
|
|
bb_width = 20
|
|
bb_height = 20
|
|
|
|
start_time = index_to_time(x_pos, spec_width, spec_duration)
|
|
end_time = index_to_time(x_pos + bb_width, spec_width, spec_duration)
|
|
low_freq = index_to_freq(y_pos, spec_height, min_freq, max_freq)
|
|
high_freq = index_to_freq(
|
|
y_pos - bb_height, spec_height, min_freq, max_freq
|
|
)
|
|
|
|
pred_nms: types.PredictionResults = {
|
|
"det_probs": np.array([1]),
|
|
"class_probs": np.array([[1]]),
|
|
"x_pos": np.array([x_pos]),
|
|
"y_pos": np.array([y_pos]),
|
|
"bb_width": np.array([bb_width]),
|
|
"bb_height": np.array([bb_height]),
|
|
"start_times": np.array([start_time]),
|
|
"end_times": np.array([end_time]),
|
|
"low_freqs": np.array([low_freq]),
|
|
"high_freqs": np.array([high_freq]),
|
|
}
|
|
|
|
params: types.FeatureExtractionParameters = {
|
|
"min_freq": min_freq,
|
|
"max_freq": max_freq,
|
|
}
|
|
|
|
features = feats.get_feats(spectrogram, pred_nms, params)
|
|
assert low_freq < high_freq
|
|
assert isinstance(features, np.ndarray)
|
|
assert features.shape == (len(pred_nms["det_probs"]), 9)
|
|
assert np.isclose(
|
|
features[0],
|
|
np.array(
|
|
[
|
|
end_time - start_time,
|
|
low_freq,
|
|
high_freq,
|
|
high_freq - low_freq,
|
|
high_freq,
|
|
max_freq,
|
|
max_freq,
|
|
max_freq,
|
|
np.nan,
|
|
]
|
|
),
|
|
equal_nan=True,
|
|
).all()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"max_power",
|
|
[
|
|
30_000,
|
|
31_000,
|
|
32_000,
|
|
33_000,
|
|
34_000,
|
|
35_000,
|
|
36_000,
|
|
37_000,
|
|
38_000,
|
|
39_000,
|
|
40_000,
|
|
],
|
|
)
|
|
def test_compute_max_power_bb(max_power: int):
|
|
"""Test compute_max_power_bb function."""
|
|
duration = 1
|
|
samplerate = 256_000
|
|
min_freq = 0
|
|
max_freq = 128_000
|
|
|
|
start_time = 0.3
|
|
end_time = 0.6
|
|
low_freq = 30_000
|
|
high_freq = 40_000
|
|
|
|
audio = np.zeros((int(duration * samplerate),))
|
|
|
|
# Add a signal during the time and frequency range of interest
|
|
audio[
|
|
int(start_time * samplerate) : int(end_time * samplerate)
|
|
] = 0.5 * librosa.tone(
|
|
max_power, sr=samplerate, duration=end_time - start_time
|
|
)
|
|
|
|
# Add a more powerful signal outside frequency range of interest
|
|
audio[
|
|
int(start_time * samplerate) : int(end_time * samplerate)
|
|
] += 2 * librosa.tone(
|
|
80_000, sr=samplerate, duration=end_time - start_time
|
|
)
|
|
|
|
params = api.get_config(
|
|
min_freq=min_freq,
|
|
max_freq=max_freq,
|
|
target_samp_rate=samplerate,
|
|
)
|
|
|
|
spec, _ = au.generate_spectrogram(
|
|
audio,
|
|
samplerate,
|
|
params,
|
|
)
|
|
|
|
x_start = int(
|
|
au.time_to_x_coords(
|
|
start_time,
|
|
samplerate,
|
|
params["fft_win_length"],
|
|
params["fft_overlap"],
|
|
)
|
|
)
|
|
|
|
x_end = int(
|
|
au.time_to_x_coords(
|
|
end_time,
|
|
samplerate,
|
|
params["fft_win_length"],
|
|
params["fft_overlap"],
|
|
)
|
|
)
|
|
|
|
num_freq_bins = spec.shape[0]
|
|
y_low = num_freq_bins - int(num_freq_bins * low_freq / max_freq)
|
|
y_high = num_freq_bins - int(num_freq_bins * high_freq / max_freq)
|
|
|
|
prediction: types.Prediction = {
|
|
"det_prob": 1,
|
|
"class_prob": np.ones((1,)),
|
|
"x_pos": x_start,
|
|
"y_pos": int(y_low),
|
|
"bb_width": int(x_end - x_start),
|
|
"bb_height": int(y_low - y_high),
|
|
"start_time": start_time,
|
|
"end_time": end_time,
|
|
"low_freq": low_freq,
|
|
"high_freq": high_freq,
|
|
}
|
|
|
|
print(prediction)
|
|
|
|
max_power_bb = feats.compute_max_power_bb(
|
|
prediction,
|
|
spec,
|
|
min_freq=min_freq,
|
|
max_freq=max_freq,
|
|
)
|
|
|
|
assert abs(max_power_bb - max_power) <= 500
|
|
|
|
|
|
def test_compute_max_power():
|
|
"""Test compute_max_power_bb function."""
|
|
duration = 3
|
|
samplerate = 16_000
|
|
min_freq = 0
|
|
max_freq = 8_000
|
|
|
|
start_time = 1
|
|
end_time = 2
|
|
low_freq = 3_000
|
|
high_freq = 4_000
|
|
max_power = 5_000
|
|
|
|
audio = np.zeros((int(duration * samplerate),))
|
|
|
|
# Add a signal during the time and frequency range of interest
|
|
audio[
|
|
int(start_time * samplerate) : int(end_time * samplerate)
|
|
] = 0.5 * librosa.tone(
|
|
3_500, sr=samplerate, duration=end_time - start_time
|
|
)
|
|
|
|
# Add a more powerful signal outside frequency range of interest
|
|
audio[
|
|
int(start_time * samplerate) : int(end_time * samplerate)
|
|
] += 2 * librosa.tone(
|
|
max_power, sr=samplerate, duration=end_time - start_time
|
|
)
|
|
|
|
params = api.get_config(
|
|
min_freq=min_freq,
|
|
max_freq=max_freq,
|
|
target_samp_rate=samplerate,
|
|
)
|
|
|
|
spec, _ = au.generate_spectrogram(
|
|
audio,
|
|
samplerate,
|
|
params,
|
|
)
|
|
|
|
x_start = int(
|
|
au.time_to_x_coords(
|
|
start_time,
|
|
samplerate,
|
|
params["fft_win_length"],
|
|
params["fft_overlap"],
|
|
)
|
|
)
|
|
|
|
x_end = int(
|
|
au.time_to_x_coords(
|
|
end_time,
|
|
samplerate,
|
|
params["fft_win_length"],
|
|
params["fft_overlap"],
|
|
)
|
|
)
|
|
|
|
num_freq_bins = spec.shape[0]
|
|
y_low = int(num_freq_bins * low_freq / max_freq)
|
|
y_high = int(num_freq_bins * high_freq / max_freq)
|
|
|
|
prediction: types.Prediction = {
|
|
"det_prob": 1,
|
|
"class_prob": np.ones((1,)),
|
|
"x_pos": x_start,
|
|
"y_pos": int(y_high),
|
|
"bb_width": int(x_end - x_start),
|
|
"bb_height": int(y_high - y_low),
|
|
"start_time": start_time,
|
|
"end_time": end_time,
|
|
"low_freq": low_freq,
|
|
"high_freq": high_freq,
|
|
}
|
|
|
|
computed_max_power = feats.compute_max_power(
|
|
prediction,
|
|
spec,
|
|
min_freq=min_freq,
|
|
max_freq=max_freq,
|
|
)
|
|
|
|
assert abs(computed_max_power - max_power) < 100
|