mecha10-nodes-object-detector 0.1.46

//! Configuration for Object Detector Node

use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Object detector node configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ObjectDetectorConfig {
    /// Path to ONNX model file
    #[serde(default = "default_model_path")]
    pub model_path: String,

    /// Model name for identification
    #[serde(default = "default_model_name")]
    pub model_name: String,

    /// Input size for YOLO model (usually 640)
    #[serde(default = "default_input_size")]
    pub input_size: u32,

    /// Confidence threshold for detections (0-1)
    #[serde(default = "default_confidence_threshold")]
    pub confidence_threshold: f32,

    /// IoU threshold for Non-Maximum Suppression (0-1)
    #[serde(default = "default_iou_threshold")]
    pub iou_threshold: f32,

    /// Maximum number of detections to return
    #[serde(default = "default_max_detections")]
    pub max_detections: usize,

    /// Frame skip - process every Nth frame (1 = process all)
    #[serde(default = "default_frame_skip")]
    pub frame_skip: u32,

    /// Start enabled by default
    #[serde(default = "default_enabled")]
    pub default_enabled: bool,

    /// Use INT8 quantization (requires quantized ONNX model)
    #[serde(default = "default_use_int8")]
    pub use_int8: bool,

    /// Max concurrent async preprocessing tasks (1 = serialize, prevents threading delay)
    #[serde(default = "default_max_async_frames")]
    pub max_async_frames: usize,

    /// Topics configuration (for framework topology)
    #[serde(default)]
    pub topics: Topics,
}

/// Topics configuration for object detector
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Topics {
    /// Topics this node publishes to
    /// Format: [{ "output": "/vision/object/detections" }]
    #[serde(default)]
    pub publishes: Vec<HashMap<String, String>>,

    /// Topics this node subscribes to
    /// Format: [{ "input": "/camera/rgb" }, { "control": "/vision/object/control" }]
    #[serde(default)]
    pub subscribes: Vec<HashMap<String, String>>,
}

impl Default for Topics {
    fn default() -> Self {
        let mut publishes = Vec::new();
        let mut subscribes = Vec::new();

        // Default publish topic
        let mut pub_map = HashMap::new();
        pub_map.insert("output".to_string(), "/vision/object/detections".to_string());
        publishes.push(pub_map);

        // Default subscribe topics
        let mut input_map = HashMap::new();
        input_map.insert("input".to_string(), "/camera/rgb".to_string());
        subscribes.push(input_map);

        let mut control_map = HashMap::new();
        control_map.insert("control".to_string(), "/vision/object/control".to_string());
        subscribes.push(control_map);

        Self { publishes, subscribes }
    }
}

impl ObjectDetectorConfig {
    /// Get input topic path from topics config
    pub fn input_topic(&self) -> String {
        self.topics
            .subscribes
            .iter()
            .find_map(|t| t.get("input"))
            .cloned()
            .unwrap_or_else(|| "/camera/rgb".to_string())
    }

    /// Get output topic path from topics config
    pub fn output_topic(&self) -> String {
        self.topics
            .publishes
            .iter()
            .find_map(|t| t.get("output"))
            .cloned()
            .unwrap_or_else(|| "/vision/object/detections".to_string())
    }

    /// Get control topic path from topics config
    pub fn control_topic(&self) -> String {
        self.topics
            .subscribes
            .iter()
            .find_map(|t| t.get("control"))
            .cloned()
            .unwrap_or_else(|| "/vision/object/control".to_string())
    }
}

fn default_model_path() -> String {
    "models/yolov8n/model.onnx".to_string()
}

fn default_model_name() -> String {
    "yolov8n".to_string()
}

fn default_input_size() -> u32 {
    640
}

fn default_confidence_threshold() -> f32 {
    0.5
}

fn default_iou_threshold() -> f32 {
    0.45
}

fn default_max_detections() -> usize {
    100
}

fn default_frame_skip() -> u32 {
    1
}

fn default_enabled() -> bool {
    false
}

fn default_use_int8() -> bool {
    false
}

fn default_max_async_frames() -> usize {
    1
}

impl Default for ObjectDetectorConfig {
    fn default() -> Self {
        Self {
            model_path: default_model_path(),
            model_name: default_model_name(),
            input_size: default_input_size(),
            confidence_threshold: default_confidence_threshold(),
            iou_threshold: default_iou_threshold(),
            max_detections: default_max_detections(),
            frame_skip: default_frame_skip(),
            default_enabled: default_enabled(),
            use_int8: default_use_int8(),
            max_async_frames: default_max_async_frames(),
            topics: Topics::default(),
        }
    }
}