batdetect2/tests/test_cli/test_process.py
2026-06-22 20:29:25 -06:00

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