11from __future__ import annotations
22
3+ import os
34from pathlib import Path
5+ from typing import TYPE_CHECKING
46
57import numpy as np
8+ from tqdm import tqdm
69
710from birdnet .acoustic .inference .core .encoding .encoding_tensor import (
811 AcousticEncodingTensor ,
912)
10- from birdnet .acoustic .inference .core .result_base import AcousticResultBase
11- from birdnet .utils .helper import get_uint_dtype
13+ from birdnet .acoustic .inference .core .result_base import (
14+ VAR_END_TIME ,
15+ VAR_INPUT ,
16+ VAR_START_TIME ,
17+ AcousticResultBase ,
18+ )
19+ from birdnet .utils .helper import (
20+ apply_speed_to_duration ,
21+ format_input_for_csv ,
22+ get_uint_dtype ,
23+ hms_centis_fast ,
24+ )
25+
26+ if TYPE_CHECKING :
27+ import pyarrow as pa
28+
29+ VAR_EMBEDDING = "embedding"
1230
1331NP_EMB_KEY = "embeddings"
1432NP_EMB_MASKED_KEY = "embeddings_masked"
@@ -53,7 +71,12 @@ def __init__(
5371
5472 @property
5573 def memory_size_MiB (self ) -> float :
56- return super ().memory_size_MiB + (
74+ """Return the total result memory usage including embeddings buffers.
75+
76+ Returns:
77+ float: Memory size in megabytes.
78+ """
79+ return super ().memory_size_mb + (
5780 (
5881 self ._embeddings .nbytes
5982 + self ._embeddings_masked .nbytes
@@ -64,21 +87,228 @@ def memory_size_MiB(self) -> float:
6487
6588 @property
6689 def embeddings (self ) -> np .ndarray :
90+ """Return the raw embedding tensor produced by the encoder.
91+
92+ Returns:
93+ np.ndarray: Embeddings with shape `(n_inputs, n_segments, emb_dim)`.
94+ """
6795 return self ._embeddings
6896
6997 @property
7098 def embeddings_masked (self ) -> np .ndarray :
99+ """Return the mask that marks relevant segments across files.
100+
101+ Returns:
102+ np.ndarray: Boolean mask of the same shape as `embeddings`.
103+ """
71104 return self ._embeddings_masked
72105
73106 @property
74- def emd_dim (self ) -> int :
107+ def emb_dim (self ) -> int :
108+ """Return the embedding dimensionality.
109+
110+ Returns:
111+ int: Number of coefficients per embedding vector.
112+ """
75113 return self ._embeddings .shape [- 1 ]
76114
77115 @property
78116 def max_n_segments (self ) -> int :
117+ """Return the maximum segment count reserved per input.
118+
119+ Returns:
120+ int: Number of overlapping windows available per file.
121+ """
79122 return self ._embeddings .shape [1 ]
80123
124+ def to_structured_array (self ) -> np .ndarray :
125+ """Convert the embeddings and timing metadata into a structured array.
126+
127+ Returns:
128+ np.ndarray: Array with fields for input path, start/end times, and embedding.
129+ """
130+ valid_mask_per_segment = ~ (self ._embeddings_masked ).all (axis = 2 )
131+ valid_file_idx , valid_seg_idx = np .where (valid_mask_per_segment )
132+ n_embeddings = len (valid_file_idx )
133+
134+ embeddings_selected = self .embeddings [valid_file_idx , valid_seg_idx ]
135+
136+ dtype = [
137+ (VAR_INPUT , self ._input_dtype ),
138+ (VAR_START_TIME , self ._input_durations .dtype ),
139+ (VAR_END_TIME , self ._input_durations .dtype ),
140+ (VAR_EMBEDDING , self ._embeddings .dtype , self .emb_dim ),
141+ ]
142+
143+ structured_array = np .empty (n_embeddings , dtype = dtype )
144+ del dtype
145+
146+ if n_embeddings == 0 :
147+ return structured_array
148+ del n_embeddings
149+
150+ sort_keys = (
151+ valid_seg_idx ,
152+ valid_file_idx ,
153+ )
154+ sort_indices = np .lexsort (sort_keys )
155+ del sort_keys
156+
157+ file_idx_flat = valid_file_idx [sort_indices ]
158+ chunk_idx_flat = valid_seg_idx [sort_indices ]
159+ emb_flat = embeddings_selected [sort_indices ]
160+ del embeddings_selected
161+ del sort_indices
162+
163+ hop_duration_s = self .hop_duration_s
164+ start_times = chunk_idx_flat .astype (self ._input_durations .dtype ) * hop_duration_s
165+ del hop_duration_s
166+ del chunk_idx_flat
167+
168+ structured_array [VAR_START_TIME ] = start_times
169+ structured_array [VAR_END_TIME ] = np .minimum (
170+ start_times
171+ + apply_speed_to_duration (self ._segment_duration_s [0 ], self ._speed [0 ]),
172+ self ._input_durations [file_idx_flat ],
173+ )
174+ del start_times
175+ structured_array [VAR_INPUT ] = self ._inputs [file_idx_flat ]
176+ del file_idx_flat
177+
178+ structured_array [VAR_EMBEDDING ] = emb_flat
179+ del emb_flat
180+
181+ return structured_array
182+
183+ def to_arrow_table (self ) -> pa .Table :
184+ """Produce a PyArrow table that serializes each embedding with timing metadata.
185+
186+ Returns:
187+ pa.Table: Table containing dictionary-encoded inputs and embeddings lists.
188+ """
189+ import pyarrow as pa
190+
191+ structured = self .to_structured_array ()
192+
193+ arrow_arrays : dict [str , pa .Array ] = {}
194+ arrow_arrays [VAR_INPUT ] = pa .array (structured [VAR_INPUT ]).dictionary_encode ()
195+ arrow_arrays [VAR_START_TIME ] = pa .array (
196+ structured [VAR_START_TIME ],
197+ type = pa .from_numpy_dtype (structured [VAR_START_TIME ].dtype ),
198+ )
199+ arrow_arrays [VAR_END_TIME ] = pa .array (
200+ structured [VAR_END_TIME ],
201+ type = pa .from_numpy_dtype (structured [VAR_END_TIME ].dtype ),
202+ )
203+
204+ embedding_element_type = pa .from_numpy_dtype (self ._embeddings .dtype )
205+ embedding_type = pa .list_ (embedding_element_type )
206+ arrow_arrays [VAR_EMBEDDING ] = pa .array (
207+ structured [VAR_EMBEDDING ].tolist (),
208+ type = embedding_type ,
209+ )
210+
211+ fields = [
212+ pa .field (VAR_INPUT , arrow_arrays [VAR_INPUT ].type , nullable = False ),
213+ pa .field (VAR_START_TIME , arrow_arrays [VAR_START_TIME ].type , nullable = False ),
214+ pa .field (VAR_END_TIME , arrow_arrays [VAR_END_TIME ].type , nullable = False ),
215+ pa .field (VAR_EMBEDDING , embedding_type , nullable = False ),
216+ ]
217+
218+ metadata : dict [bytes | str , bytes | str ] = {
219+ "segment_duration_s" : str (self ._segment_duration_s [0 ]),
220+ "overlap_duration_s" : str (self ._overlap_duration_s [0 ]),
221+ "speed" : str (self ._speed [0 ]),
222+ "n_inputs" : str (self .n_inputs ),
223+ "model_path" : str (self ._model_path [0 ]),
224+ "model_version" : str (self ._model_version [0 ]),
225+ "model_fmin" : str (self ._model_fmin [0 ]),
226+ "model_fmax" : str (self ._model_fmax [0 ]),
227+ "model_sr" : str (self ._model_sr [0 ]),
228+ "model_precision" : str (self ._model_precision [0 ]),
229+ "embedding_dim" : str (self .emb_dim ),
230+ }
231+ schema_with_metadata = pa .schema (fields , metadata = metadata )
232+ table = pa .table (arrow_arrays , schema = schema_with_metadata )
233+ return table
234+
235+ def to_csv (
236+ self ,
237+ path : os .PathLike | str ,
238+ * ,
239+ encoding : str = "utf-8" ,
240+ buffer_size_kb : int = 1024 ,
241+ silent : bool = False ,
242+ ) -> None :
243+ """Dump the structured embeddings to a CSV file for downstream analysis.
244+
245+ Args:
246+ path: File path where the CSV will be written (must end with .csv).
247+ encoding: Text encoding for the output file.
248+ buffer_size_kb: Buffer size used when writing the file.
249+ silent: Suppress progress messages when True.
250+ """
251+ if not silent :
252+ print ("Preparing CSV export..." ) # noqa: T201
253+
254+ structured = self .to_structured_array ()
255+
256+ buffer_bytes = buffer_size_kb * 1024
257+ output_path = Path (path )
258+
259+ if output_path .suffix != ".csv" :
260+ raise ValueError ("Output path must have a .csv suffix" )
261+
262+ with output_path .open ("w" , encoding = encoding , buffering = buffer_bytes ) as f :
263+ f .write (f"{ VAR_INPUT } ,{ VAR_START_TIME } ,{ VAR_END_TIME } ,{ VAR_EMBEDDING } \n " )
264+
265+ block : list [str ] = []
266+ block_size_bytes = 0
267+ total_size_bytes = 0
268+ collected_size_bytes = 0
269+ update_size_every = 1024 ** 2 * 100
270+
271+ with tqdm (
272+ total = len (structured ),
273+ desc = "Writing CSV" ,
274+ unit = "embeddings" ,
275+ disable = silent ,
276+ ) as pbar :
277+ for record in structured :
278+ line = (
279+ f"{ format_input_for_csv (record [VAR_INPUT ])} ,"
280+ f'"{ hms_centis_fast (record [VAR_START_TIME ])} ",'
281+ f'"{ hms_centis_fast (record [VAR_END_TIME ])} ",'
282+ f"{ _format_embedding_for_csv (record [VAR_EMBEDDING ])} \n "
283+ )
284+
285+ block .append (line )
286+ block_size_bytes += len (line .encode (encoding ))
287+
288+ if block_size_bytes >= buffer_bytes :
289+ f .writelines (block )
290+ block .clear ()
291+ collected_size_bytes += block_size_bytes
292+ block_size_bytes = 0
293+
294+ pbar .update (1 )
295+
296+ if collected_size_bytes >= update_size_every or pbar .n == pbar .total :
297+ total_size_bytes += collected_size_bytes
298+ collected_size_bytes = 0
299+
300+ if not silent :
301+ pbar .set_postfix ({"CSV" : f"{ total_size_bytes / 1024 ** 2 :.0f} MB" })
302+
303+ if block :
304+ f .writelines (block )
305+
81306 def unprocessable_inputs (self ) -> np .ndarray :
307+ """Return the indices of inputs that could not be processed.
308+
309+ Returns:
310+ np.ndarray: Boolean mask or indices for skipped inputs.
311+ """
82312 return self ._unprocessable_inputs
83313
84314 def _get_extra_save_data (self ) -> dict [str , np .ndarray ]:
@@ -137,9 +367,6 @@ def _input_dtype(self) -> type:
137367 # -> pointer to python string is more efficient
138368 return object
139369
140- def _format_input_for_csv (self , input_value : str ) -> str :
141- return f'"{ input_value } "'
142-
143370
144371class AcousticDataEncodingResult (AcousticEncodingResultBase ):
145372 def __init__ (
@@ -173,3 +400,9 @@ def __init__(
173400 model_precision = model_precision ,
174401 model_version = model_version ,
175402 )
403+
404+
405+ def _format_embedding_for_csv (embedding : np .ndarray , decimals : int = 6 ) -> str :
406+ fmt = f"{{:.{ decimals } f}}"
407+ formatted = "," .join (fmt .format (value ) for value in embedding )
408+ return f'"{ formatted } "'
0 commit comments