braidz_parser/
incremental_parser.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4//! This is an incremental parser for braid archives.
5
6use crate::*;
7
8/// The implementation specifies in what state we are in terms of parsing an archive.
9pub trait ParseState {}
10
11/// The archive has just been opened.
12pub struct ArchiveOpened {}
13
14/// The archive has basic information parsed.
15// The Result<> types store an error indicating why the field was not loaded.
16#[derive(Debug)]
17pub struct BasicInfoParsed {
18    pub metadata: BraidMetadata,
19    pub expected_fps: f64,
20    pub tracking_params: Option<TrackingParams>,
21    pub calibration_info: Option<CalibrationInfo>,
22    pub reconstruction_latency_hlog: Option<HistogramLog>,
23    pub reprojection_distance_hlog: Option<HistogramLog>,
24    pub cam_info: CamInfo,
25}
26
27/// The archive been completely parsed.
28#[derive(Debug)]
29pub struct FullyParsed {
30    pub metadata: BraidMetadata,
31    pub expected_fps: f64,
32    pub calibration_info: Option<CalibrationInfo>,
33    pub kalman_estimates_info: Option<KalmanEstimatesInfo>, // TODO: rename to kalman_estimates
34    pub kalman_estimates_table: Option<Vec<KalmanEstimatesRow>>,
35    pub data_association: Option<Vec<braid_types::DataAssocRow>>,
36    pub reconstruction_latency_hlog: Option<HistogramLog>,
37    pub reprojection_distance_hlog: Option<HistogramLog>,
38    pub cam_info: CamInfo,
39    pub data2d_distorted: Option<D2DInfo>,
40    /// A mapping from camera name to (width, height).
41    pub image_sizes: Option<BTreeMap<String, (usize, usize)>>,
42}
43
44impl ParseState for ArchiveOpened {}
45impl ParseState for BasicInfoParsed {}
46impl ParseState for FullyParsed {}
47
48/// An incremental parser for braid archives.
49///
50/// Initially, minimal reading from the archive is performed. As further
51/// operations on the archive proceed, the state of the parser gradually
52/// accumulates more information.
53// TODO: change this to an enum which changes its variant as it reads more.
54pub struct IncrementalParser<R: Read + Seek, S: ParseState> {
55    pub(crate) archive: zip_or_dir::ZipDirArchive<R>,
56    /// The state of parsing. Storage for stage-specific data.
57    pub(crate) state: S,
58}
59
60impl<R: Read + Seek, S: ParseState> std::fmt::Debug for IncrementalParser<R, S> {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
62        f.debug_struct("IncrementalParser")
63            .field("archive", &self.archive)
64            .finish_non_exhaustive()
65    }
66}
67
68impl IncrementalParser<BufReader<std::fs::File>, ArchiveOpened> {
69    /// Open an archive from a path.
70    ///
71    /// The archive may be a .braidz zip file for a .braid directory.
72    pub fn open<P: AsRef<std::path::Path>>(path: P) -> Result<Self, Error> {
73        let archive = zip_or_dir::ZipDirArchive::auto_from_path(path)?;
74        Ok(Self::from_archive(archive))
75    }
76
77    /// Open an archive from a directory.
78    ///
79    /// The archive must be a .braid directory.
80    pub fn open_dir<P: AsRef<std::path::Path>>(path: P) -> Result<Self, Error> {
81        let archive = zip_or_dir::ZipDirArchive::from_dir(path.as_ref().to_path_buf())?;
82        Ok(Self::from_archive(archive))
83    }
84
85    /// Open an archive from a .braidz zip file.
86    ///
87    /// The archive must be a .braidz zip file.
88    pub fn open_braidz_file<P: AsRef<std::path::Path>>(path: P) -> Result<Self, Error> {
89        let reader = BufReader::new(std::fs::File::open(&path)?);
90        let archive = zip_or_dir::ZipDirArchive::from_zip(
91            reader,
92            path.as_ref().as_os_str().to_str().unwrap().to_string(),
93        )?;
94        Ok(Self::from_archive(archive))
95    }
96}
97
98impl<R: Read + Seek> IncrementalParser<R, ArchiveOpened> {
99    /// Open an archive.
100    ///
101    /// The archive may be a .braidz zip file for a .braid directory.
102    pub fn from_archive(archive: zip_or_dir::ZipDirArchive<R>) -> Self {
103        IncrementalParser {
104            archive,
105            state: ArchiveOpened {},
106        }
107    }
108
109    /// Parse the basic data which can be quickly read from the archive.
110    pub fn parse_basics(mut self) -> Result<IncrementalParser<R, BasicInfoParsed>, Error> {
111        let mut metadata: Option<BraidMetadata> = {
112            match self.archive.open(braid_types::BRAID_METADATA_YML_FNAME) {
113                Ok(rdr) => Some(serde_yaml::from_reader(rdr)?),
114                Err(zip_or_dir::Error::FileNotFound) => None,
115                Err(e) => {
116                    return Err(Error::FileError {
117                        source: Box::new(e),
118                        filename: braid_types::BRAID_METADATA_YML_FNAME.into(),
119                        what: "opening metadata file",
120                    });
121                }
122            }
123        };
124
125        // should match:
126        //  - "unknown fps, (flydra_version 2.0.0, git_revision 0581c8fa8da4e683921480085fad21bf3b77600e, time_tzname0 CEST)"
127        //  - "100.0 fps, (top 10000, hypothesis_test_max_error 20.0)"
128        //  - "100.0 fps, (flydra_version 0.6.7, time_tzname0 CET)"
129        //  - "20.1 fps, ()"
130
131        let re_fps = regex::Regex::new(r"^(\S+) fps, \((.*)\)$").unwrap();
132
133        let re_inner =
134            regex::Regex::new(r"flydra_version (.+)\S, git_revision (\w+), time_tzname0 (.+)")
135                .unwrap();
136
137        // Parse fps and tracking parameters from textlog.
138        let mut expected_fps = f64::NAN;
139        let tracking_params: Option<TrackingParams> = {
140            let mut fname = self.archive.path_starter();
141            fname.push(braid_types::TEXTLOG_CSV_FNAME);
142            match open_maybe_gzipped(fname) {
143                Ok(rdr) => {
144                    let mut tracking_parameters = None;
145                    let textlog_rdr = csv::Reader::from_reader(rdr);
146                    for (rownum, row) in textlog_rdr.into_deserialize().early_eof_ok().enumerate() {
147                        let row: TextlogRow = row?;
148
149                        tracing::debug!(
150                            "Line in {} (row {}): {}",
151                            braid_types::TEXTLOG_CSV_FNAME,
152                            rownum,
153                            row.message
154                        );
155
156                        // TODO: fix DRY in `calc_fps_from_data()`.
157                        let line1_start = "MainBrain running at ";
158
159                        if let Some(line1_data) = row.message.strip_prefix(line1_start) {
160                            let caps = match re_fps.captures(line1_data) {
161                                Some(caps) => caps,
162                                None => return Err(Error::UnknownTextlogData),
163                            };
164                            let fps_str = caps.get(1).unwrap().as_str();
165                            let inner_str = caps.get(2).unwrap().as_str();
166                            let git_revision = match re_inner.captures(inner_str) {
167                                Some(caps2) => caps2.get(2).unwrap().as_str().to_string(),
168                                None => "unknown".to_string(),
169                            };
170
171                            if fps_str != "unknown" {
172                                expected_fps = fps_str.parse()?;
173                            }
174
175                            if metadata.is_none() {
176                                let timestamp = strand_datetime_conversion::f64_to_datetime(
177                                    row.mainbrain_timestamp,
178                                );
179
180                                let local: chrono::DateTime<chrono::Local> =
181                                    timestamp.with_timezone(&chrono::Local);
182
183                                metadata = Some(BraidMetadata {
184                                    git_revision,
185                                    original_recording_time: Some(local),
186                                    saving_program_name: "flydra".to_string(),
187                                    schema: braid_types::BRAID_SCHEMA,
188                                    save_empty_data2d: false,
189                                });
190                            }
191
192                            // No more parsing of this line. In particular, it
193                            // is not JSON.
194                            continue;
195                        }
196
197                        // parse to unstructured json
198                        let js_value_res: Result<serde_json::Value, _> =
199                            serde_json::from_str(&row.message);
200
201                        match js_value_res {
202                            Ok(mut js_value) => {
203                                if js_value
204                                    .as_object_mut()
205                                    .unwrap()
206                                    .contains_key("tracking_params")
207                                {
208                                    // If we have this key, we return an error if we
209                                    // cannot parse it.
210                                    let params_js_value = js_value["tracking_params"].take();
211                                    let tp: TrackingParams =
212                                        serde_json::from_value(params_js_value)?;
213                                    if tracking_parameters.is_some() {
214                                        return Err(Error::MultipleTrackingParameters);
215                                    }
216                                    tracking_parameters = Some(tp);
217                                }
218                            }
219                            Err(_e) => {
220                                // Cannot parse as JSON, but this is not a fatal problem.
221                                tracing::warn!(
222                                    "cannot parse message in textlog (row {rownum}) as JSON"
223                                );
224                            }
225                        }
226                    }
227
228                    tracking_parameters
229                }
230                Err(_e) => None,
231            }
232        };
233
234        let metadata = if let Some(metadata) = metadata {
235            metadata
236        } else {
237            return Err(Error::MissingMetadata {});
238        };
239
240        let calibration_info = {
241            match self.archive.open(braid_types::CALIBRATION_XML_FNAME) {
242                Ok(rdr) => {
243                    let recon: flydra_mvg::flydra_xml_support::FlydraReconstructor<f64> =
244                        serde_xml_rs::from_reader(rdr)?;
245
246                    let system =
247                        flydra_mvg::FlydraMultiCameraSystem::from_flydra_reconstructor(&recon)?;
248                    Some(CalibrationInfo {
249                        water: recon.water,
250                        cameras: system.to_system(),
251                    })
252                }
253                Err(zip_or_dir::Error::FileNotFound) => None,
254                Err(e) => {
255                    return Err(Error::FileError {
256                        source: Box::new(e),
257                        filename: braid_types::CALIBRATION_XML_FNAME.into(),
258                        what: "opening calibration file",
259                    });
260                }
261            }
262        };
263
264        let reconstruction_latency_hlog = {
265            match self
266                .archive
267                .open(braid_types::RECONSTRUCT_LATENCY_HLOG_FNAME)
268            {
269                Ok(rdr) => get_hlog(rdr).unwrap(),
270                Err(zip_or_dir::Error::FileNotFound) => None,
271                Err(e) => return Err(e.into()),
272            }
273        };
274
275        let reprojection_distance_hlog = {
276            match self.archive.open(braid_types::REPROJECTION_DIST_HLOG_FNAME) {
277                Ok(rdr) => get_hlog(rdr).unwrap(),
278                Err(zip_or_dir::Error::FileNotFound) => None,
279                Err(e) => return Err(e.into()),
280            }
281        };
282
283        let cam_info = {
284            let mut fname = self.archive.path_starter();
285            fname.push(braid_types::CAM_INFO_CSV_FNAME);
286            let rdr = open_maybe_gzipped(fname)?;
287            let caminfo_rdr = csv::Reader::from_reader(rdr);
288            let mut camn2camid = BTreeMap::new();
289            let mut camid2camn = BTreeMap::new();
290            for row in caminfo_rdr.into_deserialize().early_eof_ok() {
291                let row: CamInfoRow = row?;
292                camn2camid.insert(row.camn, row.cam_id.clone());
293                camid2camn.insert(row.cam_id, row.camn);
294            }
295            CamInfo {
296                camn2camid,
297                camid2camn,
298            }
299        };
300
301        let state = BasicInfoParsed {
302            metadata,
303            expected_fps,
304            tracking_params,
305            calibration_info,
306            reconstruction_latency_hlog,
307            reprojection_distance_hlog,
308            cam_info,
309        };
310
311        Ok(IncrementalParser {
312            archive: self.archive,
313            state,
314        })
315    }
316
317    /// Parse the entire archive.
318    pub fn parse_everything(self) -> Result<IncrementalParser<R, FullyParsed>, Error> {
319        let basics = self.parse_basics()?;
320        basics.parse_rest()
321    }
322}
323
324impl<R: Read + Seek> IncrementalParser<R, BasicInfoParsed> {
325    /// Parse the remaining aspects of the archive.
326    pub fn parse_rest(mut self) -> Result<IncrementalParser<R, FullyParsed>, Error> {
327        let basics = self.state;
328
329        let mut num_rows = 0;
330        let mut limits: Option<([u64; 2], [FlydraFloatTimestampLocal<HostClock>; 2])> = None;
331
332        let qz = {
333            // Open main 2D data.
334            let mut data_fname = self.archive.path_starter();
335            data_fname.push(braid_types::DATA2D_DISTORTED_CSV_FNAME);
336            let rdr = open_maybe_gzipped(data_fname)?;
337            let d2d_reader = csv::Reader::from_reader(rdr);
338            let mut qz = BTreeMap::new();
339
340            for row in d2d_reader.into_deserialize().early_eof_ok() {
341                num_rows += 1;
342                let row: Data2dDistortedRow = row?;
343                let entry = qz.entry(row.camn).or_insert_with(Seq2d::new);
344                if let Ok(x) = NotNan::new(row.x) {
345                    // Iff x is NaN, so is y.
346                    let y = NotNan::new(row.y).unwrap();
347                    // If 2d detection data was NaN, ignore it.
348                    entry.push(
349                        row.frame,
350                        x,
351                        y,
352                        row.timestamp,
353                        row.cam_received_timestamp.clone(),
354                    );
355                }
356                let this_frame: u64 = row.frame.try_into().unwrap();
357                let this_time = row.cam_received_timestamp;
358                if let Some((ref mut f_lim, ref mut time_lim)) = limits {
359                    f_lim[0] = std::cmp::min(f_lim[0], this_frame);
360                    f_lim[1] = std::cmp::max(f_lim[1], this_frame);
361                    time_lim[1] = this_time;
362                } else {
363                    // Initialize with the first row of data.
364                    limits = Some(([this_frame, this_frame], [this_time.clone(), this_time]));
365                }
366            }
367            qz
368        };
369
370        let data2d_distorted = limits.map(|(frame_lim, tlims)| {
371            let time_limits = [(&tlims[0]).into(), (&tlims[1]).into()];
372            D2DInfo {
373                qz,
374                frame_lim,
375                time_limits,
376                num_rows,
377            }
378        });
379
380        let data_association: Option<Vec<braid_types::DataAssocRow>> = {
381            let mut fname = self.archive.path_starter();
382            fname.push(braid_types::DATA_ASSOCIATE_CSV_FNAME);
383            let mut data_association = Vec::new();
384            match open_maybe_gzipped(fname) {
385                Ok(rdr) => {
386                    let data_assoc_reader = csv::Reader::from_reader(rdr);
387                    for row in data_assoc_reader.into_deserialize().early_eof_ok() {
388                        let row: braid_types::DataAssocRow = row?;
389                        data_association.push(row);
390                    }
391                    Some(data_association)
392                }
393                Err(e) => match e {
394                    Error::ZipOrDir {
395                        source: zip_or_dir::Error::FileNotFound,
396                    } => None,
397                    _ => {
398                        return Err(e);
399                    }
400                },
401            }
402        };
403
404        let (kalman_estimates_info, kalman_estimates_table) = {
405            let mut fname = self.archive.path_starter();
406            fname.push(braid_types::KALMAN_ESTIMATES_CSV_FNAME);
407            let mut kalman_estimates_table = Vec::new();
408            match open_maybe_gzipped(fname) {
409                Ok(rdr) => {
410                    let kest_reader = csv::Reader::from_reader(rdr);
411                    let mut trajectories = BTreeMap::new();
412                    let inf = 1.0 / 0.0;
413                    let mut xlim = [inf, -inf];
414                    let mut ylim = [inf, -inf];
415                    let mut zlim = [inf, -inf];
416                    let mut num_rows = 0;
417
418                    for row in kest_reader.into_deserialize().early_eof_ok() {
419                        let row: KalmanEstimatesRow = row?;
420                        let entry =
421                            trajectories
422                                .entry(row.obj_id)
423                                .or_insert_with(|| TrajectoryData {
424                                    // Initialize the structure with empty position vector
425                                    // and zero distance.
426                                    position: Vec::new(),
427                                    start_frame: row.frame.0,
428                                    distance: 0.0,
429                                });
430                        entry
431                            .position
432                            .push([row.x as f32, row.y as f32, row.z as f32]);
433
434                        xlim[0] = min(xlim[0], row.x);
435                        xlim[1] = max(xlim[1], row.x);
436                        ylim[0] = min(ylim[0], row.y);
437                        ylim[1] = max(ylim[1], row.y);
438                        zlim[0] = min(zlim[0], row.z);
439                        zlim[1] = max(zlim[1], row.z);
440                        num_rows += 1;
441                        kalman_estimates_table.push(row);
442                    }
443
444                    let mut total_distance: f64 = 0.0;
445                    // Loop through all individual trajectories and calculate the
446                    // distance per trajectory.
447                    for traj_data in trajectories.values_mut() {
448                        let mut previous: Option<&[f32; 3]> = None;
449                        for current in traj_data.position.iter() {
450                            if let Some(previous) = previous {
451                                let dx: f64 = (current[0] - previous[0]).into();
452                                let dy: f64 = (current[1] - previous[1]).into();
453                                let dz: f64 = (current[2] - previous[2]).into();
454                                traj_data.distance += (dx.powi(2) + dy.powi(2) + dz.powi(2)).sqrt();
455                            }
456                            previous = Some(current);
457                        }
458                        // Accumulate total distance of all trajectories.
459                        total_distance += traj_data.distance;
460                    }
461
462                    let tracking_parameters = match basics.tracking_params {
463                        Some(tp) => tp,
464                        None => {
465                            return Err(Error::MissingTrackingParameters);
466                        }
467                    };
468
469                    (
470                        Some(KalmanEstimatesInfo {
471                            xlim,
472                            ylim,
473                            zlim,
474                            trajectories,
475                            num_rows,
476                            tracking_parameters,
477                            total_distance,
478                        }),
479                        Some(kalman_estimates_table),
480                    )
481                }
482                Err(e) => match e {
483                    Error::ZipOrDir {
484                        source: zip_or_dir::Error::FileNotFound,
485                    } => (None, None),
486                    _ => {
487                        return Err(e);
488                    }
489                },
490            }
491        };
492
493        let image_sizes = if let Some(calibration_info) = basics.calibration_info.as_ref() {
494            Some(
495                calibration_info
496                    .cameras
497                    .cams_by_name()
498                    .iter()
499                    .map(|(k, v)| (k.clone(), (v.width(), v.height())))
500                    .collect(),
501            )
502        } else {
503            let mut result: BTreeMap<String, (usize, usize)> = Default::default();
504            let mut failed = false;
505            for cam_id in basics.cam_info.camid2camn.keys() {
506                let relname = format!("{}/{cam_id}.png", braid_types::IMAGES_DIRNAME);
507                match self.archive.open(relname) {
508                    Ok(mut rdr) => {
509                        let mut buf = Vec::new();
510                        rdr.read_to_end(&mut buf)?;
511                        let cur = std::io::Cursor::new(buf);
512                        let decoder = image::codecs::png::PngDecoder::new(cur)?;
513                        let (w, h) = image::ImageDecoder::dimensions(&decoder);
514                        result.insert(cam_id.clone(), (w as usize, h as usize));
515                    }
516                    Err(zip_or_dir::Error::FileNotFound) => {
517                        failed = true;
518                    }
519                    Err(e) => return Err(e.into()),
520                }
521            }
522            if !failed { Some(result) } else { None }
523        };
524
525        let cam_info = basics.cam_info;
526
527        Ok(IncrementalParser {
528            archive: self.archive,
529            state: FullyParsed {
530                metadata: basics.metadata,
531                expected_fps: basics.expected_fps,
532                calibration_info: basics.calibration_info,
533                cam_info,
534                kalman_estimates_info,
535                kalman_estimates_table,
536                data_association,
537                data2d_distorted,
538                reconstruction_latency_hlog: basics.reconstruction_latency_hlog,
539                reprojection_distance_hlog: basics.reprojection_distance_hlog,
540                image_sizes,
541            },
542        })
543    }
544
545    pub fn basic_info(&self) -> &BasicInfoParsed {
546        &self.state
547    }
548}
549
550impl<R: Read + Seek> IncrementalParser<R, FullyParsed> {
551    pub fn kalman_estimates_info(&self) -> Option<&KalmanEstimatesInfo> {
552        self.state.kalman_estimates_info.as_ref()
553    }
554}
555
556impl<R: Read + Seek, S: ParseState> IncrementalParser<R, S> {
557    /// Consume and return the raw storage archive.
558    pub fn into_inner(self) -> zip_or_dir::ZipDirArchive<R> {
559        self.archive
560    }
561
562    /// Display the path of the archive.
563    pub fn display(&self) -> std::path::Display<'_> {
564        self.archive.display()
565    }
566
567    /// Get a path-like instance for direct read access to the archive.
568    ///
569    /// You should prefer to use information already parsed from the archive
570    /// rather than resorting to this low-level function. Consider expanding the
571    /// parser to provide this information if it is not already implemented.
572    pub fn path_starter(&mut self) -> zip_or_dir::PathLike<'_, R> {
573        self.archive.path_starter()
574    }
575}
braidz_parser/incremental_parser.rs

braidz_parser/
incremental_parser.rs
