1use 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#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)]
47pub enum CameraMode {
48 PerImage,
50 Single,
52 #[default]
54 PerFolder,
55 Existing(CameraId),
57}
58
59#[derive(Clone, Debug, PartialEq)]
61pub struct ReaderOptions {
62 pub camera_model: CameraModelKind,
64 pub camera_mode: CameraMode,
66 pub camera_params: Option<Vec<f64>>,
68 pub default_focal_length_factor: f64,
70 pub max_image_size: usize,
72 pub strict_existing_camera_size: bool,
74 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#[derive(Clone, Debug, PartialEq, Default)]
94pub struct ExtractionOptions {
95 pub reader: ReaderOptions,
97 pub sift: SiftOptions,
99 pub sequential_images: bool,
101}
102
103#[derive(Clone, Debug)]
105pub struct ImageSource {
106 pub name: String,
108 pub gray: GrayImage,
110 pub exif: ExifInfo,
112}
113
114#[derive(Clone, Debug, PartialEq)]
116pub enum ImageStatus {
117 Extracted {
119 image_id: ImageId,
121 camera_id: CameraId,
123 num_features: usize,
125 },
126 AlreadyExists {
128 image_id: ImageId,
130 },
131 Failed {
133 error: String,
135 },
136}
137
138#[derive(Clone, Debug, PartialEq)]
140pub struct ImageReport {
141 pub name: String,
143 pub status: ImageStatus,
145}
146
147pub 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
153pub fn image_folder(name: &str) -> &str {
155 match name.rfind('/') {
156 Some(i) => &name[..i],
157 None => "",
158 }
159}
160
161pub 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#[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
220struct Job {
222 idx: usize,
223 image_id: ImageId,
224 camera_id: CameraId,
225}
226
227impl FeatureExtractor {
228 pub fn new() -> Self {
230 Self { backend: Arc::new(CpuSift::new()) }
231 }
232
233 pub fn with_backend(backend: Arc<dyn SiftEngine>) -> Self {
235 Self { backend }
236 }
237
238 pub fn backend(&self) -> &dyn SiftEngine {
240 self.backend.as_ref()
241 }
242
243 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 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 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 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 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 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 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
407fn 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}