mirror of
https://github.com/macaodha/batdetect2.git
synced 2026-07-07 21:00:10 +02:00
91 lines
2.3 KiB
Python
91 lines
2.3 KiB
Python
import json
|
|
import shutil
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
from click.testing import CliRunner
|
|
from soundevent.geometry import compute_bounds
|
|
|
|
from batdetect2 import BatDetect2API
|
|
from batdetect2.cli import cli
|
|
|
|
|
|
def test_cli_process_directory_merges_clip_outputs_per_recording(
|
|
tmp_path: Path,
|
|
contrib_dir: Path,
|
|
) -> None:
|
|
recording_path = contrib_dir / "jeff37" / "0166_20240531_223911.wav"
|
|
|
|
source_folder = tmp_path / "audio"
|
|
source_folder.mkdir()
|
|
shutil.copy2(
|
|
recording_path,
|
|
source_folder / "example_audio.wav",
|
|
)
|
|
|
|
destination_folder = tmp_path / "results"
|
|
destination_folder.mkdir()
|
|
|
|
api = BatDetect2API.from_checkpoint()
|
|
|
|
api_outputs = api.process_directory(
|
|
source_folder,
|
|
detection_threshold=0.3,
|
|
)
|
|
|
|
# Get all detections regardless of clip
|
|
detections = [
|
|
detection
|
|
for clip_detections in api_outputs
|
|
for detection in clip_detections.detections
|
|
]
|
|
|
|
result = CliRunner().invoke(
|
|
cli,
|
|
args=[
|
|
"process",
|
|
"directory",
|
|
str(source_folder),
|
|
str(destination_folder),
|
|
"--detection-threshold",
|
|
"0.3",
|
|
],
|
|
)
|
|
|
|
assert result.exit_code == 0
|
|
assert destination_folder.exists()
|
|
|
|
output_json = destination_folder / "example_audio.wav.json"
|
|
assert output_json.exists()
|
|
|
|
saved_detections = json.loads(output_json.read_text())
|
|
|
|
expected_annotations = Counter(
|
|
(
|
|
round(float(start_time), 4),
|
|
round(float(end_time), 4),
|
|
int(low_freq),
|
|
int(high_freq),
|
|
round(float(detection.class_scores.max()), 3),
|
|
round(float(detection.detection_score), 3),
|
|
)
|
|
for detection in detections
|
|
for start_time, low_freq, end_time, high_freq in [
|
|
compute_bounds(detection.geometry)
|
|
]
|
|
)
|
|
|
|
actual_annotations = Counter(
|
|
(
|
|
annotation["start_time"],
|
|
annotation["end_time"],
|
|
annotation["low_freq"],
|
|
annotation["high_freq"],
|
|
annotation["class_prob"],
|
|
annotation["det_prob"],
|
|
)
|
|
for annotation in saved_detections["annotation"]
|
|
)
|
|
|
|
assert actual_annotations == expected_annotations
|