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
from collections import defaultdict
from pathlib import Path
from typing import List, Literal, Sequence, TypedDict, cast
@ -93,8 +94,11 @@ class BatDetect2Formatter(OutputFormatterProtocol[FileAnnotation]):
def format(
self, predictions: Sequence[ClipDetections]
) -> List[FileAnnotation]:
merged_predictions = merge_clip_detections(predictions)
return [
self.format_prediction(prediction) for prediction in predictions
self.format_prediction(prediction)
for prediction in merged_predictions
]
def save(
@ -349,3 +353,48 @@ class BatDetect2Formatter(OutputFormatterProtocol[FileAnnotation]):
preserve_audio_tree=config.preserve_audio_tree,
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,
)