Skip to main content

yscv_cli/
source.rs

1use yscv_detect::BoundingBox;
2use yscv_recognize::Recognizer;
3use yscv_tensor::Tensor;
4use yscv_video::{
5    CameraConfig, CameraFrameSource, Frame, FrameSource, InMemoryFrameSource, resolve_camera_device,
6};
7
8use crate::config::{CliConfig, DetectTarget};
9use crate::error::AppError;
10use crate::util::embedding_from_bbox;
11
12pub fn open_camera_source(cli: &CliConfig) -> Result<CameraFrameSource, AppError> {
13    let device_index = if let Some(query) = cli.device_name_query.as_deref() {
14        let selected = resolve_camera_device(query)?;
15        println!(
16            "camera device resolved: query=`{}` -> {}: {}",
17            query, selected.index, selected.label
18        );
19        selected.index
20    } else {
21        cli.device_index
22    };
23
24    let source = CameraFrameSource::open(CameraConfig {
25        device_index,
26        width: cli.width,
27        height: cli.height,
28        fps: cli.fps,
29    })?;
30    Ok(source)
31}
32
33pub fn build_source(
34    cli: &CliConfig,
35    recognizer: &mut Recognizer,
36) -> Result<Box<dyn FrameSource>, AppError> {
37    if cli.camera {
38        let source = open_camera_source(cli)?;
39        return Ok(Box::new(source));
40    }
41
42    let frame_height = 6usize;
43    let frame_width = 6usize;
44    let frames = match cli.detect_target {
45        DetectTarget::People => vec![
46            frame_with_boxes(0, 0, frame_height, frame_width, &[(1, 1, 3, 3, 0.9)])?,
47            frame_with_boxes(1, 33_333, frame_height, frame_width, &[(1, 2, 3, 4, 0.85)])?,
48            frame_with_boxes(
49                2,
50                66_666,
51                frame_height,
52                frame_width,
53                &[(1, 3, 3, 5, 0.9), (4, 0, 6, 2, 0.88)],
54            )?,
55        ],
56        DetectTarget::Faces => vec![
57            frame_with_rgb_boxes(
58                0,
59                0,
60                frame_height,
61                frame_width,
62                &[(1, 1, 3, 3, [0.78, 0.60, 0.46])],
63                [0.10, 0.10, 0.12],
64            )?,
65            frame_with_rgb_boxes(
66                1,
67                33_333,
68                frame_height,
69                frame_width,
70                &[(1, 2, 3, 4, [0.79, 0.61, 0.47])],
71                [0.10, 0.10, 0.12],
72            )?,
73            frame_with_rgb_boxes(
74                2,
75                66_666,
76                frame_height,
77                frame_width,
78                &[
79                    (1, 3, 3, 5, [0.78, 0.60, 0.46]),
80                    (4, 0, 6, 2, [0.72, 0.55, 0.41]),
81                ],
82                [0.10, 0.10, 0.12],
83            )?,
84        ],
85    };
86
87    if recognizer.identities().is_empty() {
88        let known_embedding = embedding_from_bbox(
89            BoundingBox {
90                x1: 1.0,
91                y1: 1.0,
92                x2: 3.0,
93                y2: 3.0,
94            },
95            frame_width as f32,
96            frame_height as f32,
97        )?;
98        recognizer.enroll("alice", known_embedding)?;
99    }
100
101    Ok(Box::new(InMemoryFrameSource::new(frames)))
102}
103
104fn frame_with_boxes(
105    index: u64,
106    ts_us: u64,
107    height: usize,
108    width: usize,
109    boxes: &[(usize, usize, usize, usize, f32)],
110) -> Result<Frame, AppError> {
111    let mut data = vec![0.0f32; height * width];
112    for (x1, y1, x2, y2, value) in boxes {
113        for y in *y1..*y2 {
114            for x in *x1..*x2 {
115                data[y * width + x] = *value;
116            }
117        }
118    }
119    Ok(Frame::new(
120        index,
121        ts_us,
122        Tensor::from_vec(vec![height, width, 1], data)?,
123    )?)
124}
125
126fn frame_with_rgb_boxes(
127    index: u64,
128    ts_us: u64,
129    height: usize,
130    width: usize,
131    boxes: &[(usize, usize, usize, usize, [f32; 3])],
132    background_rgb: [f32; 3],
133) -> Result<Frame, AppError> {
134    let mut data = vec![0.0f32; height * width * 3];
135    for y in 0..height {
136        for x in 0..width {
137            let base = (y * width + x) * 3;
138            data[base] = background_rgb[0];
139            data[base + 1] = background_rgb[1];
140            data[base + 2] = background_rgb[2];
141        }
142    }
143
144    for (x1, y1, x2, y2, rgb) in boxes {
145        for y in *y1..*y2 {
146            for x in *x1..*x2 {
147                let base = (y * width + x) * 3;
148                data[base] = rgb[0];
149                data[base + 1] = rgb[1];
150                data[base + 2] = rgb[2];
151            }
152        }
153    }
154
155    Ok(Frame::new(
156        index,
157        ts_us,
158        Tensor::from_vec(vec![height, width, 3], data)?,
159    )?)
160}