From 12ccfc4ed467555e67fd51858435f2bdb4a9bdd6 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:49:08 +0100 Subject: [PATCH 1/6] Avoid np stack error for empty arrays --- src/batdetect2/outputs/formats/raw.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/batdetect2/outputs/formats/raw.py b/src/batdetect2/outputs/formats/raw.py index c4150df..9a2e6d0 100644 --- a/src/batdetect2/outputs/formats/raw.py +++ b/src/batdetect2/outputs/formats/raw.py @@ -140,7 +140,7 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): "clip_id": str(clip.uuid), } - if self.include_class_scores: + if self.include_class_scores and values["class_scores"]: class_scores = np.stack(values["class_scores"], axis=0) data_vars["class_scores"] = ( ["detection", "classes"], @@ -148,7 +148,7 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): ) coords["classes"] = ("classes", self.targets.class_names) - if self.include_features: + if self.include_features and values["features"]: features = np.stack(values["features"], axis=0) data_vars["features"] = (["detection", "feature"], features) coords["feature"] = ("feature", np.arange(num_features)) From 8475b09527ad28de4ec90eaae17ed4068192b076 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 10:41:57 +0100 Subject: [PATCH 2/6] Add progress reporting and multiprocessing --- src/batdetect2/outputs/formats/raw.py | 137 +++++++++++++++++++------- 1 file changed, 99 insertions(+), 38 deletions(-) diff --git a/src/batdetect2/outputs/formats/raw.py b/src/batdetect2/outputs/formats/raw.py index 9a2e6d0..676e517 100644 --- a/src/batdetect2/outputs/formats/raw.py +++ b/src/batdetect2/outputs/formats/raw.py @@ -1,4 +1,6 @@ +import json from collections import defaultdict +from multiprocessing import Pool from pathlib import Path from typing import List, Literal, Sequence from uuid import UUID, uuid4 @@ -8,6 +10,7 @@ import xarray as xr from loguru import logger from soundevent import data from soundevent.geometry import compute_bounds +from tqdm import tqdm from batdetect2.core import BaseConfig from batdetect2.outputs.formats.base import ( @@ -25,6 +28,8 @@ class RawOutputConfig(BaseConfig): include_class_scores: bool = True include_features: bool = True include_geometry: bool = True + n_jobs: int = 1 + show_progress: bool = False class RawFormatter(OutputFormatterProtocol[ClipDetections]): @@ -35,12 +40,19 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): include_features: bool = True, include_geometry: bool = True, parse_full_geometry: bool = False, + n_jobs: int = 1, + show_progress: bool = False, ): self.targets = targets self.include_class_scores = include_class_scores self.include_features = include_features self.include_geometry = include_geometry self.parse_full_geometry = parse_full_geometry + self.n_jobs = n_jobs + self.show_progress = show_progress + + if n_jobs < 1: + raise ValueError("n_jobs must be >= 1") def format( self, @@ -68,15 +80,40 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): def load(self, path: data.PathLike) -> List[ClipDetections]: path = Path(path) files = list(path.glob("*.nc")) - predictions: List[ClipDetections] = [] - for filepath in files: - logger.debug(f"Loading clip predictions {filepath}") - clip_data = xr.load_dataset(filepath) - prediction = self.pred_from_xr(clip_data) - predictions.append(prediction) + if self.n_jobs == 1: + return self._load_sequential(files) - return predictions + return self._load_parallel(files) + + def _load_sequential( + self, files: Sequence[data.PathLike] + ) -> List[ClipDetections]: + + iterable = files + if self.show_progress: + iterable = tqdm(files, total=len(files)) + + return [self.load_single_file(filepath) for filepath in iterable] + + def _load_parallel( + self, files: Sequence[data.PathLike] + ) -> List[ClipDetections]: + with Pool(self.n_jobs) as pool: + if not self.show_progress: + return pool.map(self.load_single_file, files) + + return list( + tqdm( + pool.imap(self.load_single_file, files), + total=len(files), + ) + ) + + def load_single_file(self, filepath: data.PathLike) -> ClipDetections: + logger.debug(f"Loading clip predictions {filepath}") + clip_data = xr.load_dataset(filepath) + return self.pred_from_xr(clip_data) def pred_to_xr( self, @@ -167,59 +204,81 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): def pred_from_xr(self, dataset: xr.Dataset) -> ClipDetections: clip_data = dataset - recording = data.Recording.model_validate_json( - clip_data.attrs["recording"] + recording = data.Recording.model_construct( + json.loads(clip_data.attrs["recording"]) ) clip_id = clip_data.clip_id.item() - clip = data.Clip( + clip = data.Clip.model_construct( recording=recording, uuid=UUID(clip_id), - start_time=clip_data.clip_start, - end_time=clip_data.clip_end, + start_time=float(clip_data.clip_start), + end_time=float(clip_data.clip_end), ) sound_events = [] - for detection in clip_data.coords["detection"]: - detection_data = clip_data.sel(detection=detection) - score = detection_data.score.item() + num_detections = len(clip_data.coords["detection"]) - if "geometry" in clip_data and self.parse_full_geometry: - geometry = data.geometry_validate( - detection_data.geometry.item() - ) + scores = clip_data.score.data + start_times = clip_data.start_time.data + end_times = clip_data.end_time.data + low_freqs = clip_data.low_freq.data + high_freqs = clip_data.high_freq.data + + top_class_scores = clip_data.top_class_score.data + top_class = clip_data.top_class.data + + num_classes = len(self.targets.class_names) + class_map = dict( + zip( + self.targets.class_names, + range(num_classes), + strict=True, + ) + ) + + geometries = None + if self.parse_full_geometry and "geometry" in clip_data: + geometries = clip_data.geometry.data + + class_scores = None + if "class_scores" in clip_data: + class_scores = clip_data.class_scores.data + + features = None + if "features" in clip_data: + features = clip_data.features.data + + for index in range(num_detections): + score = scores[index] + + if geometries is not None: + geometry = data.geometry_validate(geometries[index]) else: - start_time = detection_data.start_time.item() - end_time = detection_data.end_time.item() - low_freq = detection_data.low_freq.item() - high_freq = detection_data.high_freq.item() + start_time = start_times[index] + end_time = end_times[index] + low_freq = low_freqs[index] + high_freq = high_freqs[index] geometry = data.BoundingBox.model_construct( coordinates=[start_time, low_freq, end_time, high_freq] ) - if "class_scores" in detection_data: - class_scores = detection_data.class_scores.data + if class_scores is not None: + class_score = class_scores[index] else: - class_scores = np.zeros(len(self.targets.class_names)) - class_index = self.targets.class_names.index( - detection_data.top_class.item() - ) - class_scores[class_index] = ( - detection_data.top_class_score.item() - ) + class_score = np.zeros(num_classes) + class_index = class_map[top_class[index]] + class_score[class_index] = top_class_scores[index] - if "features" in detection_data: - features = detection_data.features.data - else: - features = np.zeros(0) + feats = features[index] if features is not None else np.zeros(0) sound_events.append( Detection( geometry=geometry, detection_score=score, - class_scores=class_scores, - features=features, + class_scores=class_score, + features=feats, ) ) @@ -236,4 +295,6 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): include_class_scores=config.include_class_scores, include_features=config.include_features, include_geometry=config.include_geometry, + n_jobs=config.n_jobs, + show_progress=config.show_progress, ) From b8f23d8b4b1e0a5d3531936e7105e127d616025e Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:49:08 +0100 Subject: [PATCH 3/6] Avoid np stack error for empty arrays --- src/batdetect2/outputs/formats/raw.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/batdetect2/outputs/formats/raw.py b/src/batdetect2/outputs/formats/raw.py index c4150df..9a2e6d0 100644 --- a/src/batdetect2/outputs/formats/raw.py +++ b/src/batdetect2/outputs/formats/raw.py @@ -140,7 +140,7 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): "clip_id": str(clip.uuid), } - if self.include_class_scores: + if self.include_class_scores and values["class_scores"]: class_scores = np.stack(values["class_scores"], axis=0) data_vars["class_scores"] = ( ["detection", "classes"], @@ -148,7 +148,7 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): ) coords["classes"] = ("classes", self.targets.class_names) - if self.include_features: + if self.include_features and values["features"]: features = np.stack(values["features"], axis=0) data_vars["features"] = (["detection", "feature"], features) coords["feature"] = ("feature", np.arange(num_features)) From 149c5b85eff698767e29ee695ab60413e53be606 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 10:41:57 +0100 Subject: [PATCH 4/6] Add progress reporting and multiprocessing --- src/batdetect2/outputs/formats/raw.py | 137 +++++++++++++++++++------- 1 file changed, 99 insertions(+), 38 deletions(-) diff --git a/src/batdetect2/outputs/formats/raw.py b/src/batdetect2/outputs/formats/raw.py index 9a2e6d0..676e517 100644 --- a/src/batdetect2/outputs/formats/raw.py +++ b/src/batdetect2/outputs/formats/raw.py @@ -1,4 +1,6 @@ +import json from collections import defaultdict +from multiprocessing import Pool from pathlib import Path from typing import List, Literal, Sequence from uuid import UUID, uuid4 @@ -8,6 +10,7 @@ import xarray as xr from loguru import logger from soundevent import data from soundevent.geometry import compute_bounds +from tqdm import tqdm from batdetect2.core import BaseConfig from batdetect2.outputs.formats.base import ( @@ -25,6 +28,8 @@ class RawOutputConfig(BaseConfig): include_class_scores: bool = True include_features: bool = True include_geometry: bool = True + n_jobs: int = 1 + show_progress: bool = False class RawFormatter(OutputFormatterProtocol[ClipDetections]): @@ -35,12 +40,19 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): include_features: bool = True, include_geometry: bool = True, parse_full_geometry: bool = False, + n_jobs: int = 1, + show_progress: bool = False, ): self.targets = targets self.include_class_scores = include_class_scores self.include_features = include_features self.include_geometry = include_geometry self.parse_full_geometry = parse_full_geometry + self.n_jobs = n_jobs + self.show_progress = show_progress + + if n_jobs < 1: + raise ValueError("n_jobs must be >= 1") def format( self, @@ -68,15 +80,40 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): def load(self, path: data.PathLike) -> List[ClipDetections]: path = Path(path) files = list(path.glob("*.nc")) - predictions: List[ClipDetections] = [] - for filepath in files: - logger.debug(f"Loading clip predictions {filepath}") - clip_data = xr.load_dataset(filepath) - prediction = self.pred_from_xr(clip_data) - predictions.append(prediction) + if self.n_jobs == 1: + return self._load_sequential(files) - return predictions + return self._load_parallel(files) + + def _load_sequential( + self, files: Sequence[data.PathLike] + ) -> List[ClipDetections]: + + iterable = files + if self.show_progress: + iterable = tqdm(files, total=len(files)) + + return [self.load_single_file(filepath) for filepath in iterable] + + def _load_parallel( + self, files: Sequence[data.PathLike] + ) -> List[ClipDetections]: + with Pool(self.n_jobs) as pool: + if not self.show_progress: + return pool.map(self.load_single_file, files) + + return list( + tqdm( + pool.imap(self.load_single_file, files), + total=len(files), + ) + ) + + def load_single_file(self, filepath: data.PathLike) -> ClipDetections: + logger.debug(f"Loading clip predictions {filepath}") + clip_data = xr.load_dataset(filepath) + return self.pred_from_xr(clip_data) def pred_to_xr( self, @@ -167,59 +204,81 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): def pred_from_xr(self, dataset: xr.Dataset) -> ClipDetections: clip_data = dataset - recording = data.Recording.model_validate_json( - clip_data.attrs["recording"] + recording = data.Recording.model_construct( + json.loads(clip_data.attrs["recording"]) ) clip_id = clip_data.clip_id.item() - clip = data.Clip( + clip = data.Clip.model_construct( recording=recording, uuid=UUID(clip_id), - start_time=clip_data.clip_start, - end_time=clip_data.clip_end, + start_time=float(clip_data.clip_start), + end_time=float(clip_data.clip_end), ) sound_events = [] - for detection in clip_data.coords["detection"]: - detection_data = clip_data.sel(detection=detection) - score = detection_data.score.item() + num_detections = len(clip_data.coords["detection"]) - if "geometry" in clip_data and self.parse_full_geometry: - geometry = data.geometry_validate( - detection_data.geometry.item() - ) + scores = clip_data.score.data + start_times = clip_data.start_time.data + end_times = clip_data.end_time.data + low_freqs = clip_data.low_freq.data + high_freqs = clip_data.high_freq.data + + top_class_scores = clip_data.top_class_score.data + top_class = clip_data.top_class.data + + num_classes = len(self.targets.class_names) + class_map = dict( + zip( + self.targets.class_names, + range(num_classes), + strict=True, + ) + ) + + geometries = None + if self.parse_full_geometry and "geometry" in clip_data: + geometries = clip_data.geometry.data + + class_scores = None + if "class_scores" in clip_data: + class_scores = clip_data.class_scores.data + + features = None + if "features" in clip_data: + features = clip_data.features.data + + for index in range(num_detections): + score = scores[index] + + if geometries is not None: + geometry = data.geometry_validate(geometries[index]) else: - start_time = detection_data.start_time.item() - end_time = detection_data.end_time.item() - low_freq = detection_data.low_freq.item() - high_freq = detection_data.high_freq.item() + start_time = start_times[index] + end_time = end_times[index] + low_freq = low_freqs[index] + high_freq = high_freqs[index] geometry = data.BoundingBox.model_construct( coordinates=[start_time, low_freq, end_time, high_freq] ) - if "class_scores" in detection_data: - class_scores = detection_data.class_scores.data + if class_scores is not None: + class_score = class_scores[index] else: - class_scores = np.zeros(len(self.targets.class_names)) - class_index = self.targets.class_names.index( - detection_data.top_class.item() - ) - class_scores[class_index] = ( - detection_data.top_class_score.item() - ) + class_score = np.zeros(num_classes) + class_index = class_map[top_class[index]] + class_score[class_index] = top_class_scores[index] - if "features" in detection_data: - features = detection_data.features.data - else: - features = np.zeros(0) + feats = features[index] if features is not None else np.zeros(0) sound_events.append( Detection( geometry=geometry, detection_score=score, - class_scores=class_scores, - features=features, + class_scores=class_score, + features=feats, ) ) @@ -236,4 +295,6 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): include_class_scores=config.include_class_scores, include_features=config.include_features, include_geometry=config.include_geometry, + n_jobs=config.n_jobs, + show_progress=config.show_progress, ) From 8f01a914d7897e4a77250801601ceb6a29e86872 Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Fri, 7 Aug 2026 19:59:32 +0100 Subject: [PATCH 5/6] test: cover raw formatter round trips --- tests/test_data/test_predictions/test_raw.py | 99 ++++++++++++++++++++ 1 file changed, 99 insertions(+) diff --git a/tests/test_data/test_predictions/test_raw.py b/tests/test_data/test_predictions/test_raw.py index f4fa1ab..374af9c 100644 --- a/tests/test_data/test_predictions/test_raw.py +++ b/tests/test_data/test_predictions/test_raw.py @@ -61,3 +61,102 @@ def test_roundtrip( ).all() assert (recovered_prediction.features == detection.features).all() assert recovered_prediction.geometry == detection.geometry + + +def test_roundtrip_recovers_recording_metadata( + sample_formatter, + create_recording, + create_clip, + sample_targets: TargetProtocol, + tmp_path: Path, +): + recording = create_recording( + tags=[data.Tag(key="source", value="test-recorder")], + duration=2, + samplerate=384_000, + time_expansion=10, + ) + clip = create_clip(recording=recording, start_time=0.25, end_time=0.75) + detection = Detection( + geometry=data.BoundingBox( + coordinates=[0.3, 45_000, 0.4, 70_000], + ), + detection_score=0.5, + class_scores=np.ones(len(sample_targets.class_names)), + features=np.ones(32), + ) + prediction = ClipDetections(clip=clip, detections=[detection]) + + path = tmp_path / "predictions" + + sample_formatter.save(predictions=[prediction], path=path) + recovered = sample_formatter.load(path=path) + + assert len(recovered) == 1 + assert recovered[0].clip.recording.model_dump(mode="json") == ( + recording.model_dump(mode="json") + ) + + +def test_roundtrip_empty_detections( + sample_formatter, + clip: data.Clip, + tmp_path: Path, +): + prediction = ClipDetections(clip=clip, detections=[]) + + path = tmp_path / "predictions" + + sample_formatter.save(predictions=[prediction], path=path) + recovered = sample_formatter.load(path=path) + + assert len(recovered) == 1 + assert recovered[0].detections == [] + assert recovered[0].clip.uuid == prediction.clip.uuid + assert recovered[0].clip.start_time == prediction.clip.start_time + assert recovered[0].clip.end_time == prediction.clip.end_time + + +def test_roundtrip_loads_with_multiprocessing( + clip: data.Clip, + sample_targets: TargetProtocol, + tmp_path: Path, +): + save_formatter = build_output_formatter( + config=RawOutputConfig(), + targets=sample_targets, + ) + load_formatter = build_output_formatter( + config=RawOutputConfig(n_jobs=2), + targets=sample_targets, + ) + predictions = [ + ClipDetections( + clip=data.Clip( + recording=clip.recording, + start_time=index, + end_time=index + 0.5, + ), + detections=[ + Detection( + geometry=data.BoundingBox( + coordinates=[index, 45_000, index + 0.1, 70_000], + ), + detection_score=0.5, + class_scores=np.ones(len(sample_targets.class_names)), + features=np.ones(32), + ) + ], + ) + for index in range(2) + ] + + path = tmp_path / "predictions" + + save_formatter.save(predictions=predictions, path=path) + recovered = load_formatter.load(path=path) + + assert len(recovered) == len(predictions) + assert {item.clip.uuid for item in recovered} == { + item.clip.uuid for item in predictions + } From 23ccf7f5d6cef9681ff6cd979eaf4aaa0bee0700 Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Sat, 8 Aug 2026 10:54:56 +0100 Subject: [PATCH 6/6] fix: validate raw recording metadata on load --- src/batdetect2/outputs/formats/raw.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/batdetect2/outputs/formats/raw.py b/src/batdetect2/outputs/formats/raw.py index 676e517..46b0ad0 100644 --- a/src/batdetect2/outputs/formats/raw.py +++ b/src/batdetect2/outputs/formats/raw.py @@ -204,7 +204,7 @@ class RawFormatter(OutputFormatterProtocol[ClipDetections]): def pred_from_xr(self, dataset: xr.Dataset) -> ClipDetections: clip_data = dataset - recording = data.Recording.model_construct( + recording = data.Recording.model_validate( json.loads(clip_data.attrs["recording"]) )