Skip to main content

cumulus3d_features/
extract.rs

1/*
2 * extract.rs
3 *
4 * Copyright (c) 2026 ParkSangWoo
5 *
6 * Author(s):
7 *
8 *      ParkSangWoo <dev.parksangwoo@gmail.com>
9 *
10 *
11 * This file is part of cumulus3d.
12 *
13 * Licensed under the Apache License, Version 2.0 (the "License");
14 * you may not use this file except in compliance with the License.
15 * You may obtain a copy of the License at
16 *
17 *      http://www.apache.org/licenses/LICENSE-2.0
18 *
19 * Unless required by applicable law or agreed to in writing, software
20 * distributed under the License is distributed on an "AS IS" BASIS,
21 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
22 * See the License for the specific language governing permissions and
23 * limitations under the License.
24 *
25 * Third-party notices: parts of the algorithms, default parameters and data
26 * formats in this file follow other open-source projects. Their copyright
27 * notices and licenses are reproduced in THIRD_PARTY_NOTICES.md.
28 *
29 * SPDX-License-Identifier: Apache-2.0
30 */
31
32//! Image list -> camera grouping -> SIFT -> incremental `FeatureStore` writes.
33
34use crate::camera_init::init_camera;
35use crate::exif_info::ExifInfo;
36use crate::gray::{self, GrayImage};
37use crate::sift::{CpuSift, SiftEngine, SiftOptions, SiftOutput};
38use cumulus3d_core::store::PosePrior;
39use cumulus3d_core::{CameraId, CameraModelKind, Error, FeatureStore, ImageId, Keypoint, Result, Vec3};
40use rayon::prelude::*;
41use std::collections::HashMap;
42use std::path::Path;
43use std::sync::Arc;
44
45/// Camera grouping mode.
46#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
47pub enum CameraMode {
48    /// A new camera per image.
49    PerImage,
50    /// All images in this call share one camera.
51    Single,
52    /// One camera per folder (parent path of the name). Sets position 0.
53    #[default]
54    PerFolder,
55    /// Attach to an existing camera id (no parameter update). Sets position p >= 1.
56    Existing(CameraId),
57}
58
59/// Image reading and camera options.
60#[derive(Clone, Debug, PartialEq)]
61pub struct ReaderOptions {
62    /// Model type of new cameras.
63    pub camera_model: CameraModelKind,
64    /// Camera grouping mode.
65    pub camera_mode: CameraMode,
66    /// If given, used as-is for the new camera parameters (pre-focal flag set).
67    pub camera_params: Option<Vec<f64>>,
68    /// Focal length when EXIF is missing = factor x max(W, H).
69    pub default_focal_length_factor: f64,
70    /// Maximum image size (downscale criterion). Effective SIFT value 3200.
71    pub max_image_size: usize,
72    /// In `Existing` mode, fail if the image size differs from the camera (if disabled, convert without error).
73    pub strict_existing_camera_size: bool,
74    /// Record EXIF GPS as the prior position.
75    pub read_pose_priors: bool,
76}
77
78impl Default for ReaderOptions {
79    fn default() -> Self {
80        Self {
81            camera_model: CameraModelKind::OpenCv,
82            camera_mode: CameraMode::PerFolder,
83            camera_params: None,
84            default_focal_length_factor: 1.2,
85            max_image_size: 3200,
86            strict_existing_camera_size: false,
87            read_pose_priors: true,
88        }
89    }
90}
91
92/// Extraction option bundle.
93#[derive(Clone, Debug, PartialEq, Default)]
94pub struct ExtractionOptions {
95    /// Image reading and camera options.
96    pub reader: ReaderOptions,
97    /// SIFT options.
98    pub sift: SiftOptions,
99    /// Process several images concurrently (memory ~ image count x pyramid).
100    pub sequential_images: bool,
101}
102
103/// In-memory input image (already decoded).
104#[derive(Clone, Debug)]
105pub struct ImageSource {
106    /// Image name in the store (relative path separated by '/').
107    pub name: String,
108    /// Grayscale image at input resolution.
109    pub gray: GrayImage,
110    /// EXIF summary of the image.
111    pub exif: ExifInfo,
112}
113
114/// Per-image processing result.
115#[derive(Clone, Debug, PartialEq)]
116pub enum ImageStatus {
117    /// Features extracted and written to the store.
118    Extracted {
119        /// Issued image id.
120        image_id: ImageId,
121        /// Assigned camera id.
122        camera_id: CameraId,
123        /// Number of features recorded.
124        num_features: usize,
125    },
126    /// Same name already exists and has features -> skipped.
127    AlreadyExists {
128        /// Existing image id.
129        image_id: ImageId,
130    },
131    /// Reading/extraction failed (processing of other images continues).
132    Failed {
133        /// Failure reason.
134        error: String,
135    },
136}
137
138/// Processing report for one image.
139#[derive(Clone, Debug, PartialEq)]
140pub struct ImageReport {
141    /// Image name in the store.
142    pub name: String,
143    /// Processing result.
144    pub status: ImageStatus,
145}
146
147/// Image list file: trim whitespace on each line, skip blank lines.
148pub fn read_image_list(path: &Path) -> Result<Vec<String>> {
149    let s = std::fs::read_to_string(path)?;
150    Ok(s.lines().map(str::trim).filter(|l| !l.is_empty()).map(str::to_string).collect())
151}
152
153/// Parent folder of the name ("camF/x.jpg" -> "camF").
154pub fn image_folder(name: &str) -> &str {
155    match name.rfind('/') {
156        Some(i) => &name[..i],
157        None => "",
158    }
159}
160
161/// Input grayscale image -> (downscale, EXIF orientation rotation) -> SIFT -> keypoints in camera-size (cam_w, cam_h) coordinates.
162pub fn extract_for_camera(
163    backend: &dyn SiftEngine,
164    gray: &GrayImage,
165    orientation: Option<u32>,
166    cam_w: u64,
167    cam_h: u64,
168    max_image_size: usize,
169    sift: &SiftOptions,
170) -> Result<(Vec<Keypoint>, SiftOutput)> {
171    let (rw, rh) = gray::limited_size(gray.width, gray.height, max_image_size);
172    let resized;
173    let img = if (rw, rh) != (gray.width, gray.height) {
174        resized = gray.resized(rw, rh);
175        &resized
176    } else {
177        gray
178    };
179    let k = gray::orientation_to_rotation(orientation);
180    let rotated;
181    let input = if k != 0 {
182        rotated = img.rotate_ccw(k);
183        &rotated
184    } else {
185        img
186    };
187    let out = backend.extract(input, sift)?;
188    let (iw, ih) = (input.width as f32, input.height as f32);
189    let sx = cam_w as f32 / rw as f32;
190    let sy = cam_h as f32 / rh as f32;
191    let kps = out
192        .features
193        .iter()
194        .map(|f| {
195            let mut kp = f.keypoint();
196            if k != 0 {
197                kp = gray::rotate_keypoint_ccw(&kp, 4 - k, iw, ih);
198            }
199            if (rw as u64, rh as u64) != (cam_w, cam_h) {
200                kp.rescale(sx, sy);
201            }
202            kp
203        })
204        .collect();
205    Ok((kps, out))
206}
207
208/// Feature extractor: backend + store writes.
209#[derive(Clone)]
210pub struct FeatureExtractor {
211    backend: Arc<dyn SiftEngine>,
212}
213
214impl Default for FeatureExtractor {
215    fn default() -> Self {
216        Self::new()
217    }
218}
219
220/// One job after camera assignment.
221struct Job {
222    idx: usize,
223    image_id: ImageId,
224    camera_id: CameraId,
225}
226
227impl FeatureExtractor {
228    /// CPU backend.
229    pub fn new() -> Self {
230        Self { backend: Arc::new(CpuSift::new()) }
231    }
232
233    /// Creates one with the given SIFT backend (e.g. GPU).
234    pub fn with_backend(backend: Arc<dyn SiftEngine>) -> Self {
235        Self { backend }
236    }
237
238    /// SIFT backend in use.
239    pub fn backend(&self) -> &dyn SiftEngine {
240        self.backend.as_ref()
241    }
242
243    /// Reads a file list (names relative to `root`) and adds them to the store. Names are processed in lexicographic order.
244    /// With `Existing(N)`, fails immediately if camera N does not exist. Per-image failures are recorded in the report.
245    pub fn extract_files(&self, store: &FeatureStore, root: &Path, names: &[String], opts: &ExtractionOptions) -> Result<Vec<ImageReport>> {
246        let mut names: Vec<String> = names.iter().map(|s| s.trim().to_string()).filter(|s| !s.is_empty()).collect();
247        names.sort();
248        names.dedup();
249        check_existing_camera(store, &opts.reader)?;
250        let mut reports: Vec<Option<ImageReport>> = vec![None; names.len()];
251        let mut todo = Vec::new();
252        for (i, n) in names.iter().enumerate() {
253            match already_done(store, n) {
254                Some(id) => reports[i] = Some(ImageReport { name: n.clone(), status: ImageStatus::AlreadyExists { image_id: id } }),
255                None => todo.push(i),
256            }
257        }
258        let decode = |i: &usize| -> std::result::Result<ImageSource, String> {
259            let n = &names[*i];
260            let path = root.join(n);
261            let bytes = std::fs::read(&path).map_err(|e| format!("{}: {e}", path.display()))?;
262            let exif = ExifInfo::from_bytes(&bytes);
263            let img = image::load_from_memory(&bytes).map_err(|e| format!("{}: 디코딩 실패: {e}", path.display()))?;
264            let gray = gray::to_gray(&img).map_err(|e| e.to_string())?;
265            Ok(ImageSource { name: n.clone(), gray, exif })
266        };
267        let decoded: Vec<std::result::Result<ImageSource, String>> =
268            if opts.sequential_images { todo.iter().map(decode).collect() } else { todo.par_iter().map(decode).collect() };
269        let mut inputs = Vec::new();
270        let mut input_slot = Vec::new();
271        for (i, d) in todo.iter().zip(decoded) {
272            match d {
273                Ok(inp) => {
274                    inputs.push(inp);
275                    input_slot.push(*i);
276                }
277                Err(e) => reports[*i] = Some(ImageReport { name: names[*i].clone(), status: ImageStatus::Failed { error: e } }),
278            }
279        }
280        let sub = self.extract_inputs(store, inputs, opts)?;
281        // extract_inputs returns results in name-sorted order. The inputs are sorted too, so the order matches.
282        for (slot, r) in input_slot.into_iter().zip(sub) {
283            reports[slot] = Some(r);
284        }
285        Ok(reports.into_iter().flatten().collect())
286    }
287
288    /// Adds decoded images to the store (for per-position calls). Issues ids in lexicographic name order.
289    pub fn extract_inputs(&self, store: &FeatureStore, mut inputs: Vec<ImageSource>, opts: &ExtractionOptions) -> Result<Vec<ImageReport>> {
290        inputs.sort_by(|a, b| a.name.cmp(&b.name));
291        let ro = &opts.reader;
292        check_existing_camera(store, ro)?;
293        let mut reports: Vec<ImageReport> =
294            inputs.iter().map(|i| ImageReport { name: i.name.clone(), status: ImageStatus::Failed { error: String::new() } }).collect();
295
296        // 1) Assign camera and image rows (sequential, deterministic ids).
297        let mut jobs = Vec::new();
298        let mut folder_cam: HashMap<String, CameraId> = HashMap::new();
299        let mut single_cam: Option<CameraId> = None;
300        for (idx, inp) in inputs.iter().enumerate() {
301            if let Some(id) = already_done(store, &inp.name) {
302                reports[idx].status = ImageStatus::AlreadyExists { image_id: id };
303                continue;
304            }
305            let (w, h) = (inp.gray.width as u64, inp.gray.height as u64);
306            let assigned: std::result::Result<(ImageId, CameraId), String> = (|| {
307                // If only the row exists without features, just refill the features.
308                if let Some(img) = store.image_by_name(&inp.name) {
309                    return Ok((img.image_id, img.camera_id));
310                }
311                let new_cam = |store: &FeatureStore| -> std::result::Result<CameraId, String> {
312                    let cam = init_camera(ro.camera_model, w, h, &inp.exif, ro.camera_params.as_deref(), ro.default_focal_length_factor)
313                        .map_err(|e| e.to_string())?;
314                    store.add_camera(cam).map_err(|e| e.to_string())
315                };
316                let check_size = |cid: CameraId| -> std::result::Result<(), String> {
317                    let c = store.camera(cid).ok_or_else(|| format!("카메라 {cid} 없음"))?;
318                    if (c.width, c.height) != (w, h) {
319                        return Err(format!("크기 불일치: 영상 {w}x{h}, 카메라 {cid} {}x{}", c.width, c.height));
320                    }
321                    Ok(())
322                };
323                let cam = match ro.camera_mode {
324                    CameraMode::PerImage => new_cam(store)?,
325                    CameraMode::Single => match single_cam {
326                        Some(c) => {
327                            check_size(c)?;
328                            c
329                        }
330                        None => {
331                            let c = new_cam(store)?;
332                            single_cam = Some(c);
333                            c
334                        }
335                    },
336                    // Design decision: group via a folder->camera map. On a sorted list this equals the
337                    // "previous camera + first-seen folder" rule.
338                    CameraMode::PerFolder => {
339                        let folder = image_folder(&inp.name).to_string();
340                        match folder_cam.get(&folder) {
341                            Some(&c) => {
342                                check_size(c)?;
343                                c
344                            }
345                            None => {
346                                let c = new_cam(store)?;
347                                folder_cam.insert(folder, c);
348                                c
349                            }
350                        }
351                    }
352                    CameraMode::Existing(n) => {
353                        if ro.strict_existing_camera_size {
354                            check_size(n)?;
355                        }
356                        n
357                    }
358                };
359                let id = store.add_image(&inp.name, cam).map_err(|e| e.to_string())?;
360                if ro.read_pose_priors {
361                    if let Some((lat, lon, alt)) = inp.exif.gps {
362                        let gravity = gray::orientation_gravity(inp.exif.orientation).map(|(gx, gy)| Vec3::new(gx, gy, 0.0));
363                        store.set_pose_prior(id, PosePrior { position: Vec3::new(lat, lon, alt), coordinate_system: 0, gravity });
364                    }
365                }
366                Ok((id, cam))
367            })();
368            match assigned {
369                Ok((image_id, camera_id)) => jobs.push(Job { idx, image_id, camera_id }),
370                Err(e) => reports[idx].status = ImageStatus::Failed { error: e },
371            }
372        }
373
374        // 2) SIFT (images in parallel) -> store.
375        let run = |job: &Job| -> (usize, ImageStatus) {
376            let inp = &inputs[job.idx];
377            let Some(cam) = store.camera(job.camera_id) else {
378                return (job.idx, ImageStatus::Failed { error: format!("카메라 {} 없음", job.camera_id) });
379            };
380            match extract_for_camera(
381                self.backend.as_ref(),
382                &inp.gray,
383                inp.exif.orientation,
384                cam.width,
385                cam.height,
386                ro.max_image_size,
387                &opts.sift,
388            ) {
389                Ok((kps, out)) => {
390                    let n = kps.len();
391                    store.set_keypoints(job.image_id, kps);
392                    store.set_descriptors(job.image_id, out.descriptors);
393                    (job.idx, ImageStatus::Extracted { image_id: job.image_id, camera_id: job.camera_id, num_features: n })
394                }
395                Err(e) => (job.idx, ImageStatus::Failed { error: e.to_string() }),
396            }
397        };
398        let results: Vec<(usize, ImageStatus)> =
399            if opts.sequential_images { jobs.iter().map(run).collect() } else { jobs.par_iter().map(run).collect() };
400        for (idx, st) in results {
401            reports[idx].status = st;
402        }
403        Ok(reports)
404    }
405}
406
407/// The id if it has a name and both keypoints and descriptors.
408fn already_done(store: &FeatureStore, name: &str) -> Option<ImageId> {
409    let id = store.image_id_by_name(name)?;
410    (store.exists_keypoints(id) && store.exists_descriptors(id)).then_some(id)
411}
412
413fn check_existing_camera(store: &FeatureStore, ro: &ReaderOptions) -> Result<Option<CameraId>> {
414    match ro.camera_mode {
415        CameraMode::Existing(n) => {
416            if store.camera(n).is_none() {
417                return Err(Error::NotFound(format!("existing_camera_id {n}")));
418            }
419            Ok(Some(n))
420        }
421        _ => Ok(None),
422    }
423}