fix: merge clip outputs for batdetect2 format

This commit is contained in:
mbsantiago 2026-06-22 20:29:25 -06:00
parent 5c300d883f
commit 3b34f467c6
4 changed files with 252 additions and 1 deletions

View File

@ -1,4 +1,5 @@
import json import json
from collections import defaultdict
from pathlib import Path from pathlib import Path
from typing import List, Literal, Sequence, TypedDict, cast from typing import List, Literal, Sequence, TypedDict, cast
@ -93,8 +94,11 @@ class BatDetect2Formatter(OutputFormatterProtocol[FileAnnotation]):
def format( def format(
self, predictions: Sequence[ClipDetections] self, predictions: Sequence[ClipDetections]
) -> List[FileAnnotation]: ) -> List[FileAnnotation]:
merged_predictions = merge_clip_detections(predictions)
return [ return [
self.format_prediction(prediction) for prediction in predictions self.format_prediction(prediction)
for prediction in merged_predictions
] ]
def save( def save(
@ -349,3 +353,48 @@ class BatDetect2Formatter(OutputFormatterProtocol[FileAnnotation]):
preserve_audio_tree=config.preserve_audio_tree, preserve_audio_tree=config.preserve_audio_tree,
include_file_path=config.include_file_path, include_file_path=config.include_file_path,
) )
def merge_clip_detections(
predictions: Sequence[ClipDetections],
) -> List[ClipDetections]:
"""Merge clip predictions into one recording-level prediction.
This intentionally discards the original clip boundaries because the
legacy BatDetect2 file format only stores recording-level detections.
"""
rec_to_clips = defaultdict(list)
rec_mapping = {}
for prediction in predictions:
recording = prediction.clip.recording
key = recording.path
rec_to_clips[key].append(prediction)
rec_mapping[key] = recording
merged_predictions = []
for rec_path, clips in rec_to_clips.items():
recording = rec_mapping[rec_path]
merged_predictions.append(
ClipDetections(
clip=data.Clip(
recording=recording,
start_time=0,
end_time=recording.duration,
),
detections=sorted(
[
detection
for clip_detections in clips
for detection in clip_detections.detections
],
key=lambda detection: (
detection.detection_score,
*compute_bounds(detection.geometry),
),
reverse=True,
),
)
)
return merged_predictions

View File

@ -0,0 +1,90 @@
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

View File

@ -0,0 +1,55 @@
from pathlib import Path
import pytest
from soundevent import data
from batdetect2.inference.batch import run_batch_inference
from batdetect2.targets import build_roi_mapping, build_targets
from batdetect2.train import load_model_from_checkpoint
from tests.utils import assert_clip_detections_equal
pytestmark = pytest.mark.slow
def test_run_batch_inference_matches_single_clip_inference(
contrib_dir: Path,
) -> None:
recording = data.Recording.from_file(
contrib_dir / "jeff37" / "0166_20240531_223911.wav"
)
clips = [
data.Clip(recording=recording, start_time=start, end_time=start + 1.0)
for start in (0.0, 1.0, 2.0)
]
model, configs = load_model_from_checkpoint()
targets = build_targets(configs.targets)
roi_mapper = build_roi_mapping(configs.targets.roi)
batched_predictions = run_batch_inference(
model,
clips,
targets=targets,
roi_mapper=roi_mapper,
batch_size=3,
num_workers=0,
)
single_predictions = [
run_batch_inference(
model,
[clip],
targets=targets,
roi_mapper=roi_mapper,
batch_size=1,
num_workers=0,
)[0]
for clip in clips
]
assert len(batched_predictions) == len(single_predictions)
for batched, single in zip(
batched_predictions,
single_predictions,
strict=True,
):
assert_clip_detections_equal(batched, single)

57
tests/utils.py Normal file
View File

@ -0,0 +1,57 @@
import numpy as np
from soundevent.geometry import compute_bounds
from batdetect2.postprocess.types import ClipDetections
def assert_clip_detections_equal(
detections: ClipDetections,
other: ClipDetections,
) -> None:
"""Assert two clip-detection objects are numerically equivalent."""
assert detections.clip.recording.path == other.clip.recording.path
assert detections.clip.start_time == other.clip.start_time
assert detections.clip.end_time == other.clip.end_time
assert len(detections.detections) == len(other.detections)
sorted_detections = sorted(
detections.detections,
key=lambda det: (
compute_bounds(det.geometry)[0],
compute_bounds(det.geometry)[1],
),
)
sorted_other = sorted(
other.detections,
key=lambda det: (
compute_bounds(det.geometry)[0],
compute_bounds(det.geometry)[1],
),
)
for det, other_det in zip(
sorted_detections,
sorted_other,
strict=True,
):
np.testing.assert_allclose(
np.array(compute_bounds(det.geometry)),
np.array(compute_bounds(other_det.geometry)),
atol=2e-2,
)
assert np.isclose(
det.detection_score,
other_det.detection_score,
atol=1e-6,
)
np.testing.assert_allclose(
det.class_scores,
other_det.class_scores,
atol=1e-6,
)
np.testing.assert_allclose(
det.features,
other_det.features,
atol=2e-6,
)