Skip to main content

flydra2/
flydra2.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! high-level multi-camera 3D tracking, ported from Flydra
5use tracing::{debug, error, info, trace};
6use tracing_futures::Instrument;
7
8use mini_arenas::MiniArenaImage;
9use serde::{Deserialize, Serialize};
10
11use std::{
12    collections::{BTreeMap, BTreeSet},
13    f64,
14    sync::{Arc, Mutex},
15};
16
17use hdrhistogram::{
18    Counter, Histogram,
19    serialization::{V2DeflateSerializer, interval_log},
20};
21
22use nalgebra::{
23    DefaultAllocator, OMatrix, OVector, Point3, RealField, Vector6,
24    allocator::Allocator,
25    dimension::{DimMin, U1, U2, U3, U6},
26};
27
28use braid_mvg::PointWorldFrame;
29
30use braid_types::{
31    CamInfoRow, CamNum, ConnectedCameraSyncState, Data2dDistortedRowF32, DataAssocRow,
32    FlydraFloatTimestampLocal, HostClock, KalmanEstimatesRow, RawCamName, SyncFno, TextlogRow,
33    TrackingParams, TriggerClockInfoRow, Triggerbox,
34};
35
36mod connected_camera_manager;
37pub use connected_camera_manager::{ConnectedCamCallback, ConnectedCamerasManager, SyncStart};
38
39mod write_data;
40pub use write_data::BraidMetadataBuilder;
41
42mod bundled_data;
43mod contiguous_stream;
44mod frame_bundler;
45
46mod new_object_test_2d;
47mod new_object_test_3d;
48
49mod flat_2d;
50mod tracking_core;
51
52mod mini_arenas;
53pub use mini_arenas::MiniArenaDebugConfig;
54
55mod model_server;
56pub use crate::model_server::{SendKalmanEstimatesRow, SendType, new_model_server};
57
58use crate::contiguous_stream::make_contiguous;
59pub use crate::frame_bundler::StreamItem;
60use crate::frame_bundler::bundle_frames;
61
62type MyFloat = braid_types::MyFloat;
63
64mod error;
65pub use error::{Error, file_error, wrap_error};
66
67pub type Result<M> = std::result::Result<M, Error>;
68
69// The first trigger pulse is labelled with this pulsenumber. Due to the
70// behavior of the triggerbox, the first pulse physically leaving the device
71// already has pulsenumber 2.
72pub const TRIGGERBOX_FIRST_PULSE: u64 = 2;
73
74pub(crate) fn generate_observation_model<R>(
75    cam: &flydra_mvg::MultiCamera<R>,
76    state: &Vector6<R>,
77    ekf_observation_covariance_pixels: f64,
78) -> Result<CameraObservationModel<R>>
79where
80    R: RealField + Copy + Default + serde::Serialize,
81{
82    let pt3d: PointWorldFrame<R> = to_world_point(state);
83    // Deals with water if needed.
84    let mat2x3 = cam.linearize_numerically_at(&pt3d, nalgebra::convert(0.001))?;
85    Ok(CameraObservationModel::new(
86        cam.clone(),
87        mat2x3,
88        ekf_observation_covariance_pixels,
89    ))
90}
91
92// We use a 6 dimensional state vector:
93// [x,y,z,xvel,yvel,zvel].
94#[derive(Debug)]
95struct CameraObservationModel<R>
96where
97    R: RealField + Copy + Default + serde::Serialize,
98{
99    cam: flydra_mvg::MultiCamera<R>,
100    observation_matrix: OMatrix<R, U2, U6>,
101    observation_matrix_transpose: OMatrix<R, U6, U2>,
102    observation_noise_covariance: OMatrix<R, U2, U2>,
103}
104
105impl<R> CameraObservationModel<R>
106where
107    R: RealField + Copy + Default + serde::Serialize,
108{
109    fn new(
110        cam: flydra_mvg::MultiCamera<R>,
111        a: OMatrix<R, U2, U3>,
112        ekf_observation_covariance_pixels: f64,
113    ) -> Self {
114        let observation_matrix = {
115            let mut o = OMatrix::<R, U2, U6>::zeros();
116            o.fixed_columns_mut::<3>(0).copy_from(&a);
117            o
118        };
119        let observation_matrix_transpose = observation_matrix.transpose();
120
121        let r = nalgebra::convert(ekf_observation_covariance_pixels);
122        let zero = nalgebra::convert(0.0);
123        let observation_noise_covariance = OMatrix::<R, U2, U2>::new(r, zero, zero, r);
124        Self {
125            cam,
126            observation_matrix,
127            observation_matrix_transpose,
128            observation_noise_covariance,
129        }
130    }
131}
132
133impl<R> adskalman::ObservationModel<R, U6, U2> for CameraObservationModel<R>
134where
135    DefaultAllocator: Allocator<U6, U6>,
136    DefaultAllocator: Allocator<U6>,
137    DefaultAllocator: Allocator<U2, U6>,
138    DefaultAllocator: Allocator<U6, U2>,
139    DefaultAllocator: Allocator<U2, U2>,
140    DefaultAllocator: Allocator<U2>,
141    U2: DimMin<U2, Output = U2>,
142    R: RealField + Copy + Default + serde::Serialize,
143{
144    fn H(&self) -> &OMatrix<R, U2, U6> {
145        &self.observation_matrix
146    }
147    fn HT(&self) -> &OMatrix<R, U6, U2> {
148        &self.observation_matrix_transpose
149    }
150    fn R(&self) -> &OMatrix<R, U2, U2> {
151        &self.observation_noise_covariance
152    }
153    fn predict_observation(&self, state: &OVector<R, U6>) -> OVector<R, U2> {
154        // TODO: update to handle water here. See tag "laksdfjasl".
155        let pt = to_world_point(state);
156        let undistored = self.cam.project_3d_to_pixel(&pt);
157        OMatrix::<R, U1, U2>::new(undistored.coords[0], undistored.coords[1]).transpose()
158        // This doesn't compile for some reason:
159        // OMatrix::<R, U2, U1>::new(undistored.coords[0], undistored.coords[1])
160    }
161}
162
163#[derive(Debug, Serialize, Deserialize)]
164pub struct ExperimentInfoRow {
165    // backward-incompatible changes here require a BRAID_SCHEMA bump (see its definition)
166    pub uuid: String,
167}
168
169#[derive(Clone, Debug, PartialEq)]
170pub struct NumberedRawUdpPoint {
171    /// the original index of the detected point
172    pub idx: u8,
173    /// the actual detected point
174    pub pt: braid_types::FlydraRawUdpPoint,
175}
176
177#[derive(Debug, Serialize, Deserialize, Clone)]
178struct TrackingParamsSaver {
179    tracking_params: braid_types::TrackingParams,
180    git_revision: String,
181}
182
183#[derive(Clone, Debug, Serialize)]
184struct SyncedFrameCount {
185    frame: SyncFno,
186}
187
188impl std::cmp::PartialEq for SyncedFrameCount {
189    fn eq(&self, other: &SyncedFrameCount) -> bool {
190        self.frame.eq(&other.frame)
191    }
192}
193
194impl std::cmp::PartialOrd for SyncedFrameCount {
195    fn partial_cmp(&self, other: &SyncedFrameCount) -> Option<std::cmp::Ordering> {
196        self.frame.partial_cmp(&other.frame)
197    }
198}
199
200#[derive(Debug, Clone)]
201pub struct TimeDataPassthrough {
202    frame: SyncFno,
203    timestamp: Option<FlydraFloatTimestampLocal<Triggerbox>>,
204}
205
206impl TimeDataPassthrough {
207    #[inline]
208    pub fn new(frame: SyncFno, timestamp: &Option<FlydraFloatTimestampLocal<Triggerbox>>) -> Self {
209        let timestamp = timestamp.clone();
210        Self { frame, timestamp }
211    }
212    /// The acquisition frame (synchronized, not raw from camera)
213    #[inline]
214    pub fn synced_frame(&self) -> SyncFno {
215        self.frame
216    }
217    /// The acquisition timestamp (synchronized, not raw from camera)
218    ///
219    /// If there is no clock model, returns None.
220    #[inline]
221    pub fn trigger_timestamp(&self) -> Option<FlydraFloatTimestampLocal<Triggerbox>> {
222        self.timestamp.clone()
223    }
224}
225
226impl std::cmp::PartialEq for TimeDataPassthrough {
227    fn eq(&self, other: &TimeDataPassthrough) -> bool {
228        let result = self.frame.eq(&other.frame);
229        if result {
230            if self.timestamp.is_none() {
231                return other.timestamp.is_none();
232            }
233
234            let ts1 = self.timestamp.clone().unwrap();
235            let ts2 = other.timestamp.clone().unwrap();
236
237            // Not sure why the timestamps may be slightly out of sync. Perhaps
238            // the time model updated in the middle of processing a
239            // frame from multiple cameras? In that case, the timestamps
240            // could indeed be slightly different from each other from the same
241            // frame.
242            if (ts1.as_f64() - ts2.as_f64()).abs() > 0.001 {
243                error!(
244                    "for frame {}: multiple timestamps {} and {} not within 1 ms",
245                    self.frame,
246                    ts1.as_f64(),
247                    ts2.as_f64()
248                );
249            }
250        }
251        result
252    }
253}
254
255fn to_world_point<R: RealField + Copy>(vec6: &OVector<R, U6>) -> PointWorldFrame<R> {
256    // TODO could we just borrow a pointer to data instead of copying it?
257    PointWorldFrame {
258        coords: Point3::new(vec6.x, vec6.y, vec6.z),
259    }
260}
261
262/// image processing results from a single camera
263#[derive(Clone, Debug, PartialEq)]
264pub struct FrameData {
265    /// camera name as kept by braid_mvg::MultiCamSystem
266    ///
267    /// This can be any UTF-8 string.
268    pub cam_name: RawCamName,
269    /// camera identification number
270    pub cam_num: CamNum,
271    /// framenumber after synchronization
272    pub synced_frame: SyncFno,
273    /// time at which hardware trigger fired
274    pub trigger_timestamp: Option<FlydraFloatTimestampLocal<Triggerbox>>,
275    /// time at which camnode got frame
276    pub cam_received_timestamp: FlydraFloatTimestampLocal<HostClock>,
277    /// timestamp from the camera
278    pub device_timestamp: Option<u64>,
279    /// frame number from the camera
280    pub block_id: Option<u64>,
281    time_delta: SyncedFrameCount,
282    tdpt: TimeDataPassthrough,
283}
284
285impl FrameData {
286    #[inline]
287    pub fn new(
288        cam_name: RawCamName,
289        cam_num: CamNum,
290        synced_frame: SyncFno,
291        trigger_timestamp: Option<FlydraFloatTimestampLocal<Triggerbox>>,
292        cam_received_timestamp: FlydraFloatTimestampLocal<HostClock>,
293        device_timestamp: Option<u64>,
294        block_id: Option<u64>,
295    ) -> Self {
296        let time_delta = Self::make_time_delta(synced_frame, trigger_timestamp.clone());
297        let tdpt = TimeDataPassthrough::new(synced_frame, &trigger_timestamp);
298        Self {
299            cam_name,
300            cam_num,
301            synced_frame,
302            trigger_timestamp,
303            cam_received_timestamp,
304            device_timestamp,
305            block_id,
306            time_delta,
307            tdpt,
308        }
309    }
310
311    #[inline]
312    fn make_time_delta(
313        synced_frame: SyncFno,
314        _trigger_timestamp: Option<FlydraFloatTimestampLocal<Triggerbox>>,
315    ) -> SyncedFrameCount {
316        SyncedFrameCount {
317            frame: synced_frame,
318        }
319    }
320}
321
322/// image processing results from a single camera on a single frame
323///
324/// This is essentially a fixed up version of the data received
325/// from the flydra UDP packet on each frame from each camera.
326#[derive(Clone, Debug, PartialEq)]
327pub struct FrameDataAndPoints {
328    pub frame_data: FrameData,
329    pub points: Vec<NumberedRawUdpPoint>,
330}
331
332impl FrameDataAndPoints {
333    fn into_save(self, save_empty_data2d: bool) -> Vec<Data2dDistortedRowF32> {
334        let frame_data = &self.frame_data;
335        let pts_to_save: Vec<Data2dDistortedRowF32> = self
336            .points
337            .iter()
338            .map(|orig| convert_to_save(frame_data, orig))
339            .collect();
340
341        let data2d_distorted: Vec<Data2dDistortedRowF32> = if !pts_to_save.is_empty() {
342            pts_to_save
343        } else if save_empty_data2d {
344            let empty_data = vec![convert_empty_to_save(frame_data)];
345            empty_data
346        } else {
347            vec![]
348        };
349        data2d_distorted
350    }
351}
352
353fn safe_u8(val: usize) -> u8 {
354    assert!(val <= u8::MAX as usize, "value out of range");
355    val as u8
356}
357
358fn convert_to_save(frame_data: &FrameData, input: &NumberedRawUdpPoint) -> Data2dDistortedRowF32 {
359    let (slope, eccentricity) = match input.pt.maybe_slope_eccentricty {
360        None => (f32::NAN, f32::NAN),
361        Some((s, e)) => (s as f32, e as f32),
362    };
363
364    Data2dDistortedRowF32 {
365        camn: frame_data.cam_num,
366        frame: frame_data.synced_frame.0 as i64,
367        timestamp: frame_data.trigger_timestamp.clone(),
368        cam_received_timestamp: frame_data.cam_received_timestamp.clone(),
369        device_timestamp: frame_data.device_timestamp,
370        block_id: frame_data.block_id,
371        x: input.pt.x0_abs as f32,
372        y: input.pt.y0_abs as f32,
373        area: input.pt.area as f32,
374        slope,
375        eccentricity,
376        frame_pt_idx: input.idx,
377        cur_val: input.pt.cur_val,
378        mean_val: input.pt.mean_val as f32,
379        sumsqf_val: input.pt.sumsqf_val as f32,
380    }
381}
382
383fn convert_empty_to_save(frame_data: &FrameData) -> Data2dDistortedRowF32 {
384    Data2dDistortedRowF32 {
385        camn: frame_data.cam_num,
386        frame: frame_data.synced_frame.0 as i64,
387        timestamp: frame_data.trigger_timestamp.clone(),
388        cam_received_timestamp: frame_data.cam_received_timestamp.clone(),
389        device_timestamp: frame_data.device_timestamp,
390        block_id: frame_data.block_id,
391        x: f32::NAN,
392        y: f32::NAN,
393        area: f32::NAN,
394        slope: f32::NAN,
395        eccentricity: f32::NAN,
396        frame_pt_idx: 0,
397        cur_val: 0,
398        mean_val: f32::NAN,
399        sumsqf_val: f32::NAN,
400    }
401}
402
403/// find all subsets of orig_set
404///
405/// translated from python version by Alex Martelli:
406/// <https://web.archive.org/web/20070331175701/http://mail.python.org/pipermail/python-list/2001-January/067815.html>
407///
408/// This is also called the power set:
409/// <http://en.wikipedia.org/wiki/Power_set>
410pub fn set_of_subsets<K, V>(orig_set: &BTreeMap<K, V>) -> BTreeSet<BTreeSet<K>>
411where
412    K: Clone + Ord,
413{
414    (0..2u32.pow(orig_set.len() as u32))
415        .map(|x| {
416            orig_set
417                .iter()
418                .enumerate()
419                .filter_map(|(i, (k, _v))| {
420                    if x & (1 << i) != 0x00 {
421                        Some(k.clone())
422                    } else {
423                        None
424                    }
425                })
426                .collect()
427        })
428        .collect()
429}
430
431#[test]
432fn test_set_of_subsets() {
433    let mut orig = BTreeMap::new();
434    orig.insert(1, 'a');
435    orig.insert(2, 'b');
436    orig.insert(3, 'c');
437
438    let result = set_of_subsets(&orig);
439
440    let expected = vec![
441        vec![],
442        vec![1],
443        vec![2],
444        vec![3],
445        vec![1, 2],
446        vec![1, 3],
447        vec![2, 3],
448        vec![1, 2, 3],
449    ];
450
451    assert_eq!(result.len(), expected.len());
452    for e in expected.into_iter() {
453        assert!(result.contains(&e.into_iter().collect::<BTreeSet<_>>()));
454    }
455}
456
457#[derive(Debug)]
458pub struct KalmanEstimateRecord {
459    pub record: KalmanEstimatesRow,
460    pub data_assoc_rows: Vec<DataAssocRow>,
461    pub mean_reproj_dist_100x: Option<u64>,
462    /// When the tracker produced this estimate. Used (together with the
463    /// trigger timestamp) for the reconstruction-latency histogram, so that
464    /// the histogram measures tracking latency rather than the delay until
465    /// the disk writer dequeues the row. `None` for backlog rows of earlier
466    /// frames saved after the fact, which are excluded from the histogram.
467    pub production_timestamp: Option<chrono::DateTime<chrono::Utc>>,
468}
469
470#[derive(Debug)]
471pub enum SaveToDiskMsg {
472    // birth?
473    KalmanEstimate(KalmanEstimateRecord),
474    // death?
475    Data2dDistorted(FrameDataAndPoints),
476    StartSavingCsv(StartSavingCsvConfig),
477    StopSavingCsv,
478    Textlog(TextlogRow),
479    TriggerClockInfo(TriggerClockInfoRow),
480    SetExperimentUuid(String),
481}
482
483/// Acts like a `csv::Writer` but buffers and orders by frame.
484///
485/// This is done to allow consumers of the kalman estimates data to iterate
486/// through the saved rows assuming that they are ordered. This assumption
487/// is easy to implicitly make, so we make it true by doing this.
488struct OrderingWriter {
489    wtr: csv::Writer<Box<dyn std::io::Write + Send>>,
490    buffer: BTreeMap<u64, Vec<KalmanEstimatesRow>>,
491}
492
493fn _test_ordering_writer_is_send() {
494    // Compile-time test to ensure OrderingWriter implements Send trait.
495    fn implements<T: Send>() {}
496    implements::<OrderingWriter>();
497}
498
499impl OrderingWriter {
500    fn new(wtr: csv::Writer<Box<dyn std::io::Write + Send>>) -> Self {
501        let buffer = BTreeMap::new();
502        Self { wtr, buffer }
503    }
504    /// Flush the writer to disk. Note this does not drain the buffer.
505    fn flush(&mut self) -> std::io::Result<()> {
506        self.wtr.flush()
507    }
508    fn serialize(&mut self, row: KalmanEstimatesRow) -> csv::Result<()> {
509        let key = row.frame.0;
510        {
511            let entry = &mut self.buffer.entry(key).or_default();
512            entry.push(row);
513        }
514
515        // Buffer up to 1000 frames, then start saving the oldest ones.
516        let buffer_size = 1000;
517        if self.buffer.len() > buffer_size {
518            let n_to_save = self.buffer.len() - buffer_size;
519            let mut to_remove: Vec<u64> = Vec::with_capacity(n_to_save);
520            {
521                for (frame, rows) in self.buffer.iter().take(n_to_save) {
522                    for row in rows.iter() {
523                        self.wtr.serialize(row)?;
524                    }
525                    to_remove.push(*frame);
526                }
527            }
528            for frame in to_remove.iter() {
529                self.buffer.remove(frame);
530            }
531        }
532        Ok(())
533    }
534}
535
536impl Drop for OrderingWriter {
537    fn drop(&mut self) {
538        // get current buffer
539        let old_buffer = std::mem::take(&mut self.buffer);
540        // drain buffer
541        for (_frame, rows) in old_buffer.into_iter() {
542            for row in rows.into_iter() {
543                self.wtr.serialize(row).expect("serialzing buffered row");
544            }
545        }
546        // flush writer
547        self.wtr.flush().expect("flush writer");
548    }
549}
550
551struct IntervalHistogram<T: Counter> {
552    histogram: Histogram<T>,
553    start_timestamp: std::time::Duration,
554    duration: std::time::Duration,
555}
556
557struct StartedHistogram<T: Counter> {
558    histogram: Histogram<T>,
559    start_timestamp: std::time::SystemTime,
560}
561
562impl<T: Counter> StartedHistogram<T> {
563    fn end(
564        self,
565        file_start_time: &std::time::SystemTime,
566        end_timestamp: std::time::SystemTime,
567    ) -> std::result::Result<IntervalHistogram<T>, std::time::SystemTimeError> {
568        let start_timestamp = self.start_timestamp.duration_since(*file_start_time)?;
569        let duration = end_timestamp.duration_since(self.start_timestamp)?;
570        Ok(IntervalHistogram {
571            histogram: self.histogram,
572            start_timestamp,
573            duration,
574        })
575    }
576}
577
578#[derive(Default)]
579struct HistogramWritingState {
580    current_store: Option<StartedHistogram<u64>>,
581    histograms: Vec<IntervalHistogram<u64>>,
582}
583
584fn save_hlog(
585    output_dirname: &std::path::Path,
586    fname: &str,
587    histograms: &[IntervalHistogram<u64>],
588    file_start_time: std::time::SystemTime,
589) {
590    // Write the reconstruction latency histograms to disk.
591    let mut log_path = output_dirname.to_path_buf();
592    log_path.push(fname);
593    log_path.set_extension("hlog");
594    let mut fd = std::fs::File::create(&log_path).expect("creating latency log file");
595
596    let mut serializer = V2DeflateSerializer::new();
597    // create a writer via a builder
598    let mut latency_log_wtr = interval_log::IntervalLogWriterBuilder::new()
599        .with_start_time(file_start_time)
600        .begin_log_with(&mut fd, &mut serializer)
601        .unwrap();
602
603    for h in histograms.iter() {
604        latency_log_wtr
605            .write_histogram(&h.histogram, h.start_timestamp, h.duration, None)
606            .unwrap();
607    }
608}
609
610fn finish_histogram(
611    hist_store: &mut Option<StartedHistogram<u64>>,
612    file_start_time: std::time::SystemTime,
613    histograms: &mut Vec<IntervalHistogram<u64>>,
614    now_system: std::time::SystemTime,
615) -> std::result::Result<(), hdrhistogram::RecordError> {
616    if let Some(hist) = hist_store.take()
617        && let Ok(h) = hist.end(&file_start_time, now_system)
618    {
619        histograms.push(h);
620    }
621    Ok(())
622}
623
624fn histogram_record(
625    value: u64,
626    hist_store: &mut Option<StartedHistogram<u64>>,
627    high: u64,
628    sigfig: u8,
629    file_start_time: std::time::SystemTime,
630    histograms: &mut Vec<IntervalHistogram<u64>>,
631    now_system: std::time::SystemTime,
632) -> std::result::Result<(), hdrhistogram::RecordError> {
633    // Create a new histogram if needed, else compute how long we have used this one.
634    let (mut hist, accum_dur) = match hist_store.take() {
635        None => {
636            // Range from 1 usec to 1 minute with 2 significant figures.
637            let hist = StartedHistogram {
638                histogram: Histogram::<u64>::new_with_bounds(1, high, sigfig).unwrap(),
639                start_timestamp: now_system,
640            };
641            (hist, None)
642        }
643        Some(hist) => {
644            let start = hist.start_timestamp;
645            (hist, now_system.duration_since(start).ok())
646        }
647    };
648
649    // Record the value in the histogram.
650    hist.histogram.record(value)?;
651
652    *hist_store = Some(hist);
653
654    if let Some(accum_dur) = accum_dur
655        && accum_dur.as_secs() >= 60
656    {
657        finish_histogram(hist_store, file_start_time, histograms, now_system)?;
658    }
659    Ok(())
660}
661
662#[derive(Debug)]
663pub struct StartSavingCsvConfig {
664    pub out_dir: std::path::PathBuf,
665    pub local: Option<chrono::DateTime<chrono::Local>>,
666    pub git_rev: String,
667    pub fps: Option<f32>,
668    pub per_cam_data: BTreeMap<RawCamName, braid_types::PerCamSaveData>,
669    pub print_stats: bool,
670    pub save_performance_histograms: bool,
671}
672
673#[derive(Debug)]
674pub struct CoordProcessorConfig {
675    pub tracking_params: TrackingParams,
676    pub save_empty_data2d: bool,
677    pub ignore_latency: bool,
678    pub mini_arena_debug_cfg: Option<mini_arenas::MiniArenaDebugConfig>,
679    pub write_buffer_size_num_messages: usize,
680}
681
682/// A [tokio::sync::mpsc::Sender] which cannot be cloned.
683///
684/// This prevents accidentally keeping the receiver open because there can only
685/// be the one sender.
686///
687/// (Note that this is not a hard guarantee. A clone could be made by upgrading
688/// a `WeakSender` to a full-fledged `Sender`. Potentially new Downgraded and
689/// Upgraded types could be invented which would eliminate this possibility.)
690#[derive(Debug)]
691pub struct SingletonSender<T>(tokio::sync::mpsc::Sender<T>);
692
693impl<T> SingletonSender<T> {
694    pub async fn send(
695        &self,
696        msg: T,
697    ) -> std::result::Result<(), tokio::sync::mpsc::error::SendError<T>> {
698        self.0.send(msg).await
699    }
700
701    pub fn downgrade(&self) -> tokio::sync::mpsc::WeakSender<T> {
702        self.0.downgrade()
703    }
704}
705
706// TODO note: currently, clones of `braidz_write_tx` keep the writing task alive
707// (and thus prevent it from being dropped and saving files). We should consider
708// refactoring this so that mostly only Weak<Sender<_>> copies of `braidz_write_tx`
709// are kept and thus that the sender will drop when needed. The alternative (or
710// addition) is to have a message which will close the writer's files, as is
711// done with `SaveToDiskMsg::StopSavingCsv`.
712#[derive(Debug)]
713pub struct CoordProcessor {
714    pub cam_manager: ConnectedCamerasManager,
715    pub recon: Option<flydra_mvg::FlydraMultiCameraSystem<MyFloat>>, // TODO? keep reference
716    /// Channel to send messages to the writing thread.
717    pub braidz_write_tx: SingletonSender<SaveToDiskMsg>,
718    pub writer_join_handle: tokio::task::JoinHandle<Result<()>>,
719    model_servers: Vec<tokio::sync::mpsc::Sender<(SendType, TimeDataPassthrough)>>,
720    tracking_params: Arc<TrackingParams>,
721    /// Images of the "mini arenas" in use.
722    ///
723    /// One per camera when we have calibrations to do tracking. Empty
724    /// otherwise.
725    mini_arena_images: std::collections::BTreeMap<String, MiniArenaImage>,
726    /// A vector of model collections, one per "mini arena".
727    ///
728    /// This is behind `Option<>` for reasons I do not remember.
729    model_collections: Option<
730        Vec<crate::tracking_core::ModelCollection<crate::tracking_core::CollectionFrameDone>>,
731    >,
732    next_obj_id: Arc<Mutex<u32>>,
733}
734
735impl CoordProcessor {
736    #[tracing::instrument(level = "debug", skip_all)]
737    pub fn new(
738        cfg: CoordProcessorConfig,
739        cam_manager: ConnectedCamerasManager,
740        recon: Option<flydra_mvg::FlydraMultiCameraSystem<MyFloat>>,
741        metadata_builder: BraidMetadataBuilder,
742    ) -> Result<Self> {
743        let CoordProcessorConfig {
744            tracking_params,
745            save_empty_data2d,
746            ignore_latency,
747            mini_arena_debug_cfg,
748            write_buffer_size_num_messages,
749        } = cfg;
750
751        trace!("CoordProcessor using {:?}", recon);
752
753        let recon2 = recon.clone();
754
755        info!("using TrackingParams {:?}", tracking_params);
756
757        let mini_arena_images = mini_arenas::build_mini_arena_images(
758            recon.as_ref(),
759            &tracking_params.mini_arena_config,
760            mini_arena_debug_cfg.as_ref(),
761        )?;
762
763        let tracking_params: Arc<TrackingParams> = Arc::from(tracking_params);
764        let tracking_params2 = tracking_params.clone();
765        let cam_manager2 = cam_manager.clone();
766
767        let (braidz_write_tx, braidz_write_rx) =
768            tokio::sync::mpsc::channel(write_buffer_size_num_messages);
769
770        let writer_join_handle = tokio::task::spawn_blocking(move || {
771            match write_data::writer_task_main(
772                braidz_write_rx,
773                cam_manager2,
774                recon2,
775                tracking_params2,
776                save_empty_data2d,
777                metadata_builder,
778                ignore_latency,
779            ) {
780                Ok(()) => Ok(()),
781                Err(err) => {
782                    use std::error::Error;
783                    error!("Braidz writer task failed: {}", err);
784                    let mut outer = &err as &(dyn Error + 'static);
785                    while let Some(source) = outer.source() {
786                        error!("Cause: {source}");
787                        outer = source;
788                    }
789                    Err(err)
790                }
791            }
792        });
793
794        Ok(Self {
795            cam_manager,
796            recon,
797            braidz_write_tx: SingletonSender(braidz_write_tx),
798            writer_join_handle,
799            tracking_params,
800            model_servers: vec![],
801            model_collections: None,
802            mini_arena_images,
803            next_obj_id: Arc::new(Mutex::new(0)),
804        })
805    }
806
807    fn new_model_collections(
808        &self,
809        recon: &flydra_mvg::FlydraMultiCameraSystem<MyFloat>,
810        fps: f32,
811    ) -> Vec<crate::tracking_core::ModelCollection<crate::tracking_core::CollectionFrameDone>> {
812        self.tracking_params
813            .mini_arena_config
814            .iter_locators()
815            .map(|mini_arena_loc| {
816                let mini_arena_idx =
817                    mini_arenas::MiniArenaIndex::new(mini_arena_loc.idx().unwrap());
818                crate::tracking_core::initialize_model_collection(
819                    self.tracking_params.clone(),
820                    recon.clone(),
821                    fps,
822                    self.cam_manager.clone(),
823                    mini_arena_idx,
824                )
825            })
826            .collect()
827    }
828
829    pub fn add_listener(
830        &mut self,
831        model_server: tokio::sync::mpsc::Sender<(SendType, TimeDataPassthrough)>,
832    ) {
833        self.model_servers.push(model_server);
834    }
835
836    /// Consume the CoordProcessor and the input stream.
837    ///
838    /// Returns a future that completes when done. This is basically the "main
839    /// loop". It is async, though, and yields many times throughout this
840    /// execution.
841    ///
842    /// Upon completion, returns a [std::thread::JoinHandle] from a spawned
843    /// writing thread. To ensure data is completely saved, this should be
844    /// driven to completion before ending the process.
845    #[tracing::instrument(level = "debug", skip_all)]
846    pub async fn consume_stream<S>(
847        mut self,
848        frame_data_rx: S,
849        expected_framerate: Option<f32>,
850    ) -> Result<tokio::task::JoinHandle<Result<()>>>
851    where
852        S: 'static + Send + futures::stream::Stream<Item = StreamItem>,
853    {
854        let mut prev_frame = SyncFno(0);
855        use futures::stream::StreamExt;
856
857        // As first step, save raw incoming data. The raw data is saved by
858        // cloning each packet and sending this to the writing task. A new
859        // `Stream<Item = StreamItem>` is returned which simply moves the items
860        // from the original stream.
861        let stream1 = Box::pin(frame_data_rx.then(|si: StreamItem| async {
862            match &si {
863                StreamItem::EOF => {}
864                StreamItem::Packet(fdp) => {
865                    if fdp.frame_data.synced_frame.0 == u64::MAX {
866                        // We have seen a bug after making a contiguous stream
867                        // (see below) in which the frame number is `u64::MAX`.
868                        // This checks if this obviously wrong frame number is
869                        // introduced after the present location or before. In
870                        // any case, if we are getting frame numbers like this,
871                        // clearly we cannot track anymore, so panicing here
872                        // only raises the issue slightly earlier.
873                        panic!("Impossible frame number with frame data {fdp:?}");
874                    }
875
876                    self.braidz_write_tx
877                        .send(SaveToDiskMsg::Data2dDistorted(fdp.clone()))
878                        .await
879                        .unwrap();
880                }
881            }
882            si
883        }));
884
885        // This clones the `Arc` but the inner camera manager remains not
886        // cloned.
887        let ccm = self.cam_manager.clone();
888
889        info!("Starting model collection and frame bundler.");
890
891        // Start the model collection.
892
893        if let Some(ref recon) = self.recon {
894            let fps = expected_framerate.expect("expected_framerate must be set");
895            self.model_collections = Some(self.new_model_collections(recon, fps));
896            let dummy_time = TimeDataPassthrough {
897                frame: SyncFno(0),
898                timestamp: None,
899            };
900            // send calibration here
901            let mut flydra_xml_new: Vec<u8> = Vec::new();
902            recon
903                .to_flydra_xml(&mut flydra_xml_new)
904                .expect("to_flydra_xml");
905            let flydra_xml_str = std::str::from_utf8(&flydra_xml_new).unwrap();
906
907            for ms in self.model_servers.iter() {
908                ms.send((
909                    SendType::CalibrationFlydraXml(flydra_xml_str.to_string()),
910                    dummy_time.clone(),
911                ))
912                .await
913                .expect("send calibration");
914            }
915        }
916
917        // Start the frame bundler.
918
919        // This function takes a stream and returns a stream. In the returned
920        // stream, it has bundled the camera-by-camera data into all-cam data.
921        // Note that this can drop data that is out-of-order, which is why we
922        // must save the incoming data before here.
923        let bundled =
924            bundle_frames(stream1, ccm.clone()).instrument(tracing::info_span!("bundle_frames"));
925
926        // Ensure that there are no skipped frames.
927        let mut contiguous_stream =
928            make_contiguous(bundled).instrument(tracing::info_span!("contiguous"));
929
930        let mut mini_arena_assignment_debug = std::env::var_os("DEBUG_MINI_ARENAS")
931            .map(|fname| mini_arenas::MiniArenaAssignmentDebug::new(fname).unwrap());
932
933        // In this inner loop, we handle each incoming datum. We spend the vast majority
934        // of the runtime in this loop.
935        while let Some(bundle) = contiguous_stream.next().await {
936            assert!(
937                bundle.frame() >= prev_frame,
938                "Frame number decreasing? The previously received frame was {}, but now have {}",
939                prev_frame,
940                bundle.frame()
941            );
942            prev_frame = bundle.frame();
943
944            // Undistort incoming points and assign to mini arenas.
945            let undistorted = if let Some(recon) = &self.recon {
946                bundle.undistort_and_split_to_mini_arenas(
947                    recon,
948                    &self.mini_arena_images,
949                    &self.tracking_params.mini_arena_config,
950                )
951            } else {
952                continue;
953            };
954
955            if let Some(dbg) = mini_arena_assignment_debug.as_mut() {
956                // This uses blocking IO. It should be rewritten to use async IO.
957                dbg.write_frame(&undistorted)?;
958            }
959
960            if let Some(mcs) = &self.model_collections {
961                debug_assert_eq!(undistorted.per_mini_arena.len(), mcs.len());
962            }
963
964            // TODO: split processing across arenas into multiple threads.
965            if let Some(model_collections) = self.model_collections.take() {
966                // Across all arenas, predict motion (Kalman prediction step).
967                let model_collections = model_collections
968                    .into_iter()
969                    .map(|mc| mc.predict_motion())
970                    .collect::<Vec<_>>();
971
972                let tdpt = &undistorted.tdpt;
973
974                // ---------------------------------
975                // ---------------------------------
976                // ---------------------------------
977
978                // Across all arenas, compute likelihood of each observation.
979                let model_collections = model_collections
980                    .into_iter()
981                    .zip(undistorted.per_mini_arena.iter())
982                    .map(|(mc, arena_bundle)| mc.compute_observation_likes(tdpt, arena_bundle))
983                    .collect::<Vec<_>>();
984
985                // Across all arenas, perform data association
986                let model_collections_and_unused_observations = model_collections
987                    .into_iter()
988                    .zip(undistorted.per_mini_arena.into_iter())
989                    .map(|(mc, arena_bundle)| {
990                        mc.solve_data_association_and_update(tdpt, arena_bundle)
991                    })
992                    .collect::<Vec<_>>();
993
994                // ---------------------------------
995                // ---------------------------------
996                // ---------------------------------
997
998                // create new and delete old objects
999                let (model_collections, combined) = model_collections_and_unused_observations
1000                    .into_iter()
1001                    .map(|(mc, unused)| {
1002                        let (mc, send_msgs, save_msgs) =
1003                            mc.births_and_deaths(tdpt, unused, || self.next_obj_id_func());
1004                        (mc, (send_msgs, save_msgs))
1005                    })
1006                    .unzip::<_, _, Vec<_>, Vec<_>>();
1007
1008                for (send_msgs, save_msgs) in combined.into_iter() {
1009                    for msg in save_msgs.into_iter() {
1010                        self.braidz_write_tx.send(msg).await.unwrap();
1011                    }
1012                    for ms in self.model_servers.iter() {
1013                        for msg in send_msgs.iter() {
1014                            ms.send(msg.clone()).await.unwrap();
1015                        }
1016                    }
1017                }
1018
1019                self.model_collections = Some(model_collections);
1020            }
1021        }
1022        debug!("consume_stream future done");
1023
1024        Ok(self.writer_join_handle)
1025    }
1026
1027    fn next_obj_id_func(&self) -> u32 {
1028        let mut guard = self.next_obj_id.lock().unwrap();
1029        let val: u32 = *guard;
1030        *guard += 1;
1031        val
1032    }
1033}
1034
1035#[derive(Debug, Clone)]
1036pub(crate) struct CamAndDist {
1037    pub(crate) raw_cam_name: RawCamName,
1038    /// The reprojection distance of the undistorted pixels.
1039    pub(crate) reproj_dist: MyFloat,
1040}
1041
1042pub(crate) struct HypothesisTestResult {
1043    pub(crate) coords: Point3<MyFloat>,
1044    pub(crate) cams_and_reproj_dist: Vec<CamAndDist>,
1045}
1046
1047#[test]
1048fn test_csv_nan() {
1049    // test https://github.com/BurntSushi/rust-csv/issues/153
1050
1051    let save_row_data = Data2dDistortedRowF32 {
1052        camn: CamNum(1),
1053        frame: 2,
1054        timestamp: None,
1055        cam_received_timestamp: FlydraFloatTimestampLocal::from_dt(&chrono::Local::now()),
1056        device_timestamp: None,
1057        block_id: None,
1058        x: f32::NAN,
1059        y: f32::NAN,
1060        area: 1.0,
1061        slope: 2.0,
1062        eccentricity: 3.0,
1063        frame_pt_idx: 4,
1064        cur_val: 5,
1065        mean_val: 6.0,
1066        sumsqf_val: 7.0,
1067    };
1068
1069    let mut csv_buf = Vec::<u8>::new();
1070
1071    {
1072        let mut wtr = csv::Writer::from_writer(&mut csv_buf);
1073        wtr.serialize(&save_row_data).unwrap();
1074    }
1075
1076    println!("{}", std::str::from_utf8(&csv_buf).unwrap());
1077
1078    {
1079        let rdr = csv::Reader::from_reader(csv_buf.as_slice());
1080        let mut count = 0;
1081        for row in rdr.into_deserialize() {
1082            let row: braid_types::Data2dDistortedRow = row.unwrap();
1083            count += 1;
1084            assert!(row.x.is_nan());
1085            assert!(row.y.is_nan());
1086            assert!(!row.area.is_nan());
1087            assert_eq!(row.area, 1.0);
1088        }
1089        assert_eq!(count, 1);
1090    }
1091}