similari 0.22.6

Machine learning framework for building object trackers and similarity search engines
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

use crate::prelude::SortTrack;
use crossbeam::channel::{Receiver, Sender};
use log::debug;
use pyo3::prelude::*;
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

pub type BatchRecords<T> = HashMap<u64, Vec<T>>;
pub type SceneTracks = (u64, Vec<SortTrack>);

#[derive(Debug, Clone)]
pub struct PredictionBatchRequest<T> {
    batch: BatchRecords<T>,
    sender: Sender<SceneTracks>,
    batch_size: Arc<Mutex<usize>>,
}

#[pyclass]
#[derive(Clone, Debug)]
pub struct PredictionBatchResult {
    receiver: Receiver<SceneTracks>,
    batch_size: Arc<Mutex<usize>>,
}

#[pymethods]
impl PredictionBatchResult {
    pub fn ready(&self) -> bool {
        !self.receiver.is_empty()
    }

    #[pyo3(name = "get", text_signature = "($self)")]
    fn get_py(&self) -> SceneTracks {
        let gil = Python::acquire_gil();
        let py = gil.python();
        py.allow_threads(|| self.get())
    }

    pub fn batch_size(&self) -> usize {
        *self.batch_size.lock().unwrap()
    }
}

impl PredictionBatchResult {
    pub fn get(&self) -> SceneTracks {
        self.receiver
            .recv()
            .expect("Receiver must always receive batch computation result")
    }
}

impl<T> PredictionBatchRequest<T> {
    pub fn get_sender(&self) -> Sender<SceneTracks> {
        self.sender.clone()
    }

    #[allow(dead_code)]
    pub(crate) fn send(&self, res: SceneTracks) -> bool {
        let res = self.sender.send(res);
        if let Err(e) = res {
            debug!(
                "Error occurred when sending results to the batch result object. Error is: {:?}",
                e
            );
            false
        } else {
            true
        }
    }

    pub fn batch_size(&self) -> usize {
        *self.batch_size.lock().unwrap()
    }

    pub fn add(&mut self, scene_id: u64, elt: T) {
        let vec = self.batch.get_mut(&scene_id);
        if let Some(vec) = vec {
            vec.push(elt);
        } else {
            self.batch.insert(scene_id, vec![elt]);
        }
        let mut batch_size = self.batch_size.lock().unwrap();
        *batch_size = self.batch.len();
    }

    pub fn new() -> (Self, PredictionBatchResult) {
        let (sender, receiver) = crossbeam::channel::bounded(1);
        let batch_size = Arc::new(Mutex::new(0));
        (
            Self {
                batch: BatchRecords::default(),
                sender,
                batch_size: batch_size.clone(),
            },
            PredictionBatchResult {
                receiver,
                batch_size,
            },
        )
    }

    pub fn get_batch(&self) -> &BatchRecords<T> {
        &self.batch
    }
}

#[cfg(test)]
mod tests {
    use crate::prelude::Universal2DBox;
    use crate::trackers::batch::PredictionBatchRequest;

    #[test]
    fn test() {
        let (mut request, result) = PredictionBatchRequest::<Universal2DBox>::new();
        request.add(0, Universal2DBox::new(0.0, 0.0, Some(0.5), 1.0, 5.0));
        request.add(0, Universal2DBox::new(5.0, 5.0, Some(0.0), 1.5, 10.0));
        request.add(1, Universal2DBox::new(0.0, 0.0, Some(1.0), 0.7, 5.1));
        let batch = request.get_batch();
        drop(batch);
        assert_eq!(result.batch_size(), 2);

        assert!(request.send((0, vec![])));
        assert_eq!(result.ready(), true);
        let res = result.get();
        assert_eq!(res.0, 0);
        assert!(res.1.is_empty());
        drop(result);
        assert!(!request.send((0, vec![])));
    }
}