407 lines
15 KiB
Rust
407 lines
15 KiB
Rust
// use anyhow::{Context, Result};
|
|
// use chrono::Utc;
|
|
// use clap::Parser;
|
|
// use log::{debug, error, info, warn};
|
|
// use opencv::{core, highgui, imgcodecs, imgproc, prelude::*, videoio};
|
|
// use std::path::{Path, PathBuf};
|
|
// use std::sync::{Arc, Mutex};
|
|
// use std::time::{Duration, Instant};
|
|
// use tokio::runtime::Runtime;
|
|
|
|
// use meteor_detect::camera::frame_buffer::Frame;
|
|
// use meteor_detect::detection::{
|
|
// AggregationStrategy, BrightnessDetector, BrightnessDetectorParams,
|
|
// CamsDetector, CamsDetectorParams, DetectionResult, DetectorConfig,
|
|
// DetectorFactory, MeteorDetector, PipelineConfig
|
|
// };
|
|
|
|
// /// Command line arguments for detector demo
|
|
// #[derive(Parser, Debug)]
|
|
// #[clap(author, version, about = "Meteor detection demo")]
|
|
// struct Args {
|
|
// /// Input video file path or camera device number
|
|
// #[clap(short, long)]
|
|
// input: String,
|
|
|
|
// /// Output directory for feature images
|
|
// #[clap(short, long, default_value = "output")]
|
|
// output_dir: PathBuf,
|
|
|
|
// /// Brightness threshold for meteor detection
|
|
// #[clap(long, default_value = "30")]
|
|
// brightness_threshold: u8,
|
|
|
|
// /// Std-to-avg ratio threshold for meteor detection
|
|
// #[clap(long, default_value = "1.5")]
|
|
// std_ratio_threshold: f32,
|
|
|
|
// /// Minimum pixel count to trigger detection
|
|
// #[clap(long, default_value = "10")]
|
|
// min_pixel_count: u32,
|
|
|
|
// /// Minimum trajectory length
|
|
// #[clap(long, default_value = "5")]
|
|
// min_trajectory_length: u32,
|
|
|
|
// /// Save all feature images, not just detections
|
|
// #[clap(long)]
|
|
// save_all: bool,
|
|
|
|
// /// Display frames in real-time
|
|
// #[clap(short, long)]
|
|
// display: bool,
|
|
|
|
// /// Batch size (number of frames to process)
|
|
// #[clap(short, long, default_value = "256")]
|
|
// batch_size: usize,
|
|
|
|
// /// Detector type (cams, brightness, both)
|
|
// #[clap(short, long, default_value = "cams")]
|
|
// detector: String,
|
|
|
|
// /// Minimum brightness delta for brightness detector
|
|
// #[clap(long, default_value = "30.0")]
|
|
// min_brightness_delta: f32,
|
|
|
|
// /// Minimum pixel change for brightness detector
|
|
// #[clap(long, default_value = "10")]
|
|
// min_pixel_change: u32,
|
|
|
|
// /// Minimum consecutive frames for brightness detector
|
|
// #[clap(long, default_value = "3")]
|
|
// min_frames: u32,
|
|
|
|
// /// Sensitivity for brightness detector
|
|
// #[clap(long, default_value = "0.7")]
|
|
// sensitivity: f32,
|
|
|
|
// /// Run detection in parallel
|
|
// #[clap(long)]
|
|
// parallel: bool,
|
|
|
|
// /// Result aggregation strategy (any, all, majority)
|
|
// #[clap(long, default_value = "any")]
|
|
// aggregation: String,
|
|
// }
|
|
|
|
// fn main() -> Result<()> {
|
|
// // Set up logging
|
|
// env_logger::init();
|
|
|
|
// // Parse command line arguments
|
|
// let args = Args::parse();
|
|
|
|
// // Create output directory if it doesn't exist
|
|
// std::fs::create_dir_all(&args.output_dir)?;
|
|
|
|
// // Set up detector parameters
|
|
// let cams_params = CamsDetectorParams {
|
|
// brightness_threshold: args.brightness_threshold,
|
|
// std_to_avg_ratio_threshold: args.std_ratio_threshold,
|
|
// min_pixel_count: args.min_pixel_count,
|
|
// min_trajectory_length: args.min_trajectory_length,
|
|
// save_all_feature_images: args.save_all,
|
|
// output_dir: args.output_dir.clone(),
|
|
// file_prefix: "meteor".to_string(),
|
|
// id: "cams-demo".to_string(),
|
|
// };
|
|
|
|
// let brightness_params = BrightnessDetectorParams {
|
|
// min_brightness_delta: args.min_brightness_delta,
|
|
// min_pixel_change: args.min_pixel_change,
|
|
// min_frames: args.min_frames,
|
|
// sensitivity: args.sensitivity,
|
|
// id: "brightness-demo".to_string(),
|
|
// };
|
|
|
|
// // Create detectors based on command line argument
|
|
// let mut detectors: Vec<Box<dyn MeteorDetector>> = Vec::new();
|
|
|
|
// match args.detector.as_str() {
|
|
// "cams" => {
|
|
// detectors.push(Box::new(CamsDetector::with_params(cams_params)));
|
|
// info!("Using CAMS detector");
|
|
// },
|
|
// "brightness" => {
|
|
// detectors.push(Box::new(BrightnessDetector::with_params(brightness_params)));
|
|
// info!("Using brightness detector");
|
|
// },
|
|
// "both" => {
|
|
// detectors.push(Box::new(CamsDetector::with_params(cams_params)));
|
|
// detectors.push(Box::new(BrightnessDetector::with_params(brightness_params)));
|
|
// info!("Using both CAMS and brightness detectors");
|
|
// },
|
|
// _ => {
|
|
// return Err(anyhow::anyhow!("Unknown detector type: {}", args.detector));
|
|
// }
|
|
// }
|
|
|
|
// // Set up detection pipeline config if running in parallel
|
|
// let aggregation_strategy = match args.aggregation.as_str() {
|
|
// "any" => AggregationStrategy::Any,
|
|
// "all" => AggregationStrategy::All,
|
|
// "majority" => AggregationStrategy::Majority,
|
|
// _ => AggregationStrategy::Any,
|
|
// };
|
|
|
|
// // Open video source
|
|
// let mut cap = if args.input.chars().all(char::is_numeric) {
|
|
// // Input is a camera device number
|
|
// let device_id = args.input.parse::<i32>().unwrap_or(0);
|
|
// videoio::VideoCapture::new(device_id, videoio::CAP_ANY)?
|
|
// } else {
|
|
// // Input is a video file
|
|
// videoio::VideoCapture::from_file(&args.input, videoio::CAP_ANY)?
|
|
// };
|
|
|
|
// // Check if video source is opened
|
|
// if !cap.is_opened()? {
|
|
// return Err(anyhow::anyhow!("Failed to open video source: {}", args.input));
|
|
// }
|
|
|
|
// // Get video properties
|
|
// let fps = cap.get(videoio::CAP_PROP_FPS)?;
|
|
// let frame_width = cap.get(videoio::CAP_PROP_FRAME_WIDTH)? as i32;
|
|
// let frame_height = cap.get(videoio::CAP_PROP_FRAME_HEIGHT)? as i32;
|
|
// let frame_count = cap.get(videoio::CAP_PROP_FRAME_COUNT)? as i64;
|
|
|
|
// info!(
|
|
// "Video source: {}x{}, {:.2} fps, {} frames",
|
|
// frame_width, frame_height, fps, frame_count
|
|
// );
|
|
|
|
// // Create window if display is enabled
|
|
// if args.display {
|
|
// highgui::named_window("Video", highgui::WINDOW_NORMAL)?;
|
|
// highgui::resize_window("Video", 800, 600)?;
|
|
// }
|
|
|
|
// // Initialize frame counter and runtime
|
|
// let mut frame_idx: u64 = 0;
|
|
// let mut runtime = if args.parallel {
|
|
// Some(Runtime::new()?)
|
|
// } else {
|
|
// None
|
|
// };
|
|
|
|
// // Start time
|
|
// let start_time = Instant::now();
|
|
|
|
// // Process frames
|
|
// let mut frame = core::Mat::default();
|
|
// let mut detections = 0;
|
|
|
|
// // Convert detectors to thread-safe Arc<Mutex<Box<dyn MeteorDetector>>> for parallel mode
|
|
// let shared_detectors: Vec<Arc<Mutex<Box<dyn MeteorDetector>>>> = detectors
|
|
// .into_iter()
|
|
// .map(|d| Arc::new(Mutex::new(d)))
|
|
// .collect();
|
|
|
|
// while cap.read(&mut frame)? {
|
|
// // Skip empty frames
|
|
// if frame.empty() {
|
|
// warn!("Empty frame received, skipping");
|
|
// continue;
|
|
// }
|
|
|
|
// // Process frame with detector(s)
|
|
// let results = if args.parallel && runtime.is_some() {
|
|
// // Process in parallel using tokio runtime
|
|
// let runtime = runtime.as_ref().unwrap();
|
|
// runtime.block_on(async {
|
|
// // Create a future for each detector
|
|
// let mut handles = Vec::with_capacity(shared_detectors.len());
|
|
|
|
// for detector in &shared_detectors {
|
|
// let frame = frame.clone();
|
|
// let detector = detector.clone();
|
|
|
|
// let handle = tokio::spawn(async move {
|
|
// let mut detector = detector.lock().unwrap();
|
|
// match detector.process_frame(&frame, frame_idx) {
|
|
// Ok(result) => Some(result),
|
|
// Err(e) => {
|
|
// error!("Error processing frame with detector {}: {}",
|
|
// detector.get_id(), e);
|
|
// None
|
|
// }
|
|
// }
|
|
// });
|
|
|
|
// handles.push(handle);
|
|
// }
|
|
|
|
// // Wait for all detectors to complete
|
|
// let results = futures::future::join_all(handles).await;
|
|
|
|
// // Collect results
|
|
// results.into_iter()
|
|
// .filter_map(|r| r.ok().flatten())
|
|
// .collect::<Vec<_>>()
|
|
// })
|
|
// } else {
|
|
// // Process sequentially
|
|
// let mut results = Vec::new();
|
|
|
|
// for detector in &shared_detectors {
|
|
// let mut detector = detector.lock().unwrap();
|
|
// match detector.process_frame(&frame, frame_idx) {
|
|
// Ok(result) => {
|
|
// results.push(result);
|
|
// }
|
|
// Err(e) => {
|
|
// error!("Error processing frame with detector {}: {}",
|
|
// detector.get_id(), e);
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
// results
|
|
// };
|
|
|
|
// // Aggregate results based on strategy
|
|
// let detected = match aggregation_strategy {
|
|
// AggregationStrategy::Any => results.iter().any(|r| r.detected),
|
|
// AggregationStrategy::All => results.iter().all(|r| r.detected) && !results.is_empty(),
|
|
// AggregationStrategy::Majority => {
|
|
// let count = results.iter().filter(|r| r.detected).count();
|
|
// count > results.len() / 2 && !results.is_empty()
|
|
// },
|
|
// _ => results.iter().any(|r| r.detected),
|
|
// };
|
|
|
|
// // Find result with highest confidence
|
|
// let result = if detected {
|
|
// results.iter()
|
|
// .filter(|r| r.detected)
|
|
// .max_by(|a, b| a.confidence.partial_cmp(&b.confidence).unwrap())
|
|
// .cloned()
|
|
// } else {
|
|
// results.first().cloned()
|
|
// };
|
|
|
|
// if let Some(result) = result {
|
|
// if result.detected {
|
|
// detections += 1;
|
|
// info!(
|
|
// "Meteor detected at frame {}: confidence={:.2}, pixels={}, detector={}",
|
|
// frame_idx, result.confidence, result.pixel_change,
|
|
// result.detector_id.as_deref().unwrap_or("unknown")
|
|
// );
|
|
|
|
// // Draw bounding box if we have one
|
|
// if let Some(bbox) = result.bounding_box {
|
|
// let rect = core::Rect::new(
|
|
// bbox.0 as i32,
|
|
// bbox.1 as i32,
|
|
// bbox.2 as i32,
|
|
// bbox.3 as i32,
|
|
// );
|
|
// let mut display_frame = frame.clone();
|
|
// imgproc::rectangle(
|
|
// &mut display_frame,
|
|
// rect,
|
|
// core::Scalar::new(0.0, 255.0, 0.0, 0.0),
|
|
// 2,
|
|
// imgproc::LINE_8,
|
|
// 0,
|
|
// )?;
|
|
|
|
// // Save detection frame
|
|
// let detection_path = args
|
|
// .output_dir
|
|
// .join(format!("detection_{}_frame.jpg", frame_idx));
|
|
// imgcodecs::imwrite(
|
|
// detection_path.to_str().unwrap(),
|
|
// &display_frame,
|
|
// &core::Vector::new(),
|
|
// )?;
|
|
// }
|
|
// }
|
|
|
|
// // Display progress
|
|
// if frame_idx % 100 == 0 {
|
|
// let elapsed = start_time.elapsed();
|
|
// let fps = frame_idx as f64 / elapsed.as_secs_f64();
|
|
// debug!(
|
|
// "Processed {} frames ({:.2} fps), {} detections",
|
|
// frame_idx,
|
|
// fps,
|
|
// detections,
|
|
// );
|
|
// }
|
|
|
|
// // Display frame if enabled
|
|
// if args.display {
|
|
// // Draw text with frame information
|
|
// let mut display_frame = frame.clone();
|
|
// let text = format!(
|
|
// "Frame: {} | Detections: {} | {}",
|
|
// frame_idx,
|
|
// detections,
|
|
// if detected {
|
|
// format!("METEOR DETECTED - Confidence: {:.2}", result.confidence)
|
|
// } else {
|
|
// "No detection".to_string()
|
|
// }
|
|
// );
|
|
// imgproc::put_text(
|
|
// &mut display_frame,
|
|
// &text,
|
|
// core::Point::new(20, 40),
|
|
// imgproc::FONT_HERSHEY_SIMPLEX,
|
|
// 0.7,
|
|
// core::Scalar::new(0.0, 255.0, 0.0, 0.0),
|
|
// 2,
|
|
// imgproc::LINE_8,
|
|
// false,
|
|
// )?;
|
|
|
|
// // Show frame
|
|
// highgui::imshow("Video", &display_frame)?;
|
|
|
|
// // Wait for key press (30ms delay for video playback)
|
|
// let key = highgui::wait_key(30)?;
|
|
// if key > 0 && key != 255 {
|
|
// // 'q' or ESC key pressed
|
|
// if key == 113 || key == 27 {
|
|
// info!("User requested exit");
|
|
// break;
|
|
// }
|
|
// }
|
|
// }
|
|
// }
|
|
|
|
// frame_idx += 1;
|
|
// }
|
|
|
|
// // Close video display window if enabled
|
|
// if args.display {
|
|
// highgui::destroy_all_windows()?;
|
|
// }
|
|
|
|
// // Calculate overall statistics
|
|
// let elapsed = start_time.elapsed();
|
|
// let overall_fps = frame_idx as f64 / elapsed.as_secs_f64();
|
|
|
|
// info!(
|
|
// "Processing complete: {} frames, {} detections, {:.2} fps average",
|
|
// frame_idx, detections, overall_fps
|
|
// );
|
|
|
|
// // Display detection mode info
|
|
// info!(
|
|
// "Detection mode: {}, {}",
|
|
// if args.parallel { "parallel" } else { "sequential" },
|
|
// match aggregation_strategy {
|
|
// AggregationStrategy::Any => "any detector",
|
|
// AggregationStrategy::All => "all detectors",
|
|
// AggregationStrategy::Majority => "majority of detectors",
|
|
// _ => "custom strategy",
|
|
// }
|
|
// );
|
|
|
|
// Ok(())
|
|
// }
|
|
|
|
fn main(){} |