rustytorch/demos/object-detector-shared/src/lib.rs

//! Shared IPC types for the Object Detector demo
//!
//! This crate defines the data structures shared between the Rust backend
//! and the TypeScript frontend for the YOLO-style object detection demo.

use serde::{Deserialize, Serialize};

/// Supported YOLO model variants
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[allow(non_camel_case_types)]
pub enum DetectorModel {
    /// YOLOv8 Nano (smallest, fastest)
    #[serde(rename = "yolov8_n")]
    YOLOv8_N,
    /// YOLOv8 Small
    #[serde(rename = "yolov8_s")]
    YOLOv8_S,
    /// YOLOv8 Medium
    #[serde(rename = "yolov8_m")]
    YOLOv8_M,
    /// YOLOv8 Large
    #[serde(rename = "yolov8_l")]
    YOLOv8_L,
    /// YOLOv8 XLarge (largest, most accurate)
    #[serde(rename = "yolov8_x")]
    YOLOv8_X,
}

impl DetectorModel {
    /// Get human-readable name
    pub fn display_name(&self) -> &'static str {
        match self {
            Self::YOLOv8_N => "YOLOv8-Nano",
            Self::YOLOv8_S => "YOLOv8-Small",
            Self::YOLOv8_M => "YOLOv8-Medium",
            Self::YOLOv8_L => "YOLOv8-Large",
            Self::YOLOv8_X => "YOLOv8-XLarge",
        }
    }

    /// Get approximate parameter count
    pub fn param_count(&self) -> usize {
        match self {
            Self::YOLOv8_N => 3_200_000,  // 3.2M
            Self::YOLOv8_S => 11_200_000, // 11.2M
            Self::YOLOv8_M => 25_900_000, // 25.9M
            Self::YOLOv8_L => 43_700_000, // 43.7M
            Self::YOLOv8_X => 68_200_000, // 68.2M
        }
    }
}

/// Bounding box in normalized coordinates (0.0-1.0)
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct BoundingBox {
    /// X coordinate of top-left corner (normalized 0-1)
    pub x: f32,
    /// Y coordinate of top-left corner (normalized 0-1)
    pub y: f32,
    /// Width (normalized 0-1)
    pub width: f32,
    /// Height (normalized 0-1)
    pub height: f32,
}

impl BoundingBox {
    /// Calculate area of bounding box
    pub fn area(&self) -> f32 {
        self.width * self.height
    }

    /// Calculate center point (cx, cy)
    pub fn center(&self) -> (f32, f32) {
        (self.x + self.width / 2.0, self.y + self.height / 2.0)
    }

    /// Calculate intersection area with another box
    pub fn intersection(&self, other: &Self) -> f32 {
        let x1 = self.x.max(other.x);
        let y1 = self.y.max(other.y);
        let x2 = (self.x + self.width).min(other.x + other.width);
        let y2 = (self.y + self.height).min(other.y + other.height);

        let width = (x2 - x1).max(0.0);
        let height = (y2 - y1).max(0.0);

        width * height
    }

    /// Calculate Intersection over Union (IoU) with another box
    pub fn iou(&self, other: &Self) -> f32 {
        let intersection = self.intersection(other);
        let union = self.area() + other.area() - intersection;

        if union > 0.0 {
            intersection / union
        } else {
            0.0
        }
    }
}

/// A single detection result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Detection {
    /// Bounding box
    pub bbox: BoundingBox,
    /// Class ID (0-79 for COCO)
    pub class_id: usize,
    /// Human-readable class name
    pub class_name: String,
    /// Confidence score (0.0-1.0)
    pub confidence: f32,
    /// Color for visualization (hex string)
    pub color: String,
}

/// Result of object detection
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectionResult {
    /// List of detections
    pub detections: Vec<Detection>,
    /// Inference time in milliseconds
    pub inference_time_ms: f64,
    /// Input image width
    pub image_width: usize,
    /// Input image height
    pub image_height: usize,
}

/// Configuration for the object detector
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectorConfig {
    /// Model variant to use
    pub model: DetectorModel,
    /// Confidence threshold (0.0-1.0)
    pub confidence_threshold: f32,
    /// Non-maximum suppression threshold (0.0-1.0)
    pub nms_threshold: f32,
    /// Maximum number of detections to return
    pub max_detections: usize,
    /// Whether to use GPU if available
    pub use_gpu: bool,
}

impl Default for DetectorConfig {
    fn default() -> Self {
        Self {
            model: DetectorModel::YOLOv8_N,
            confidence_threshold: 0.25,
            nms_threshold: 0.45,
            max_detections: 100,
            use_gpu: true,
        }
    }
}

/// Status of the detector service
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DetectorStatus {
    /// Whether the model is loaded
    pub initialized: bool,
    /// Current model (if loaded)
    pub model: Option<String>,
    /// Compute device being used
    pub device: String,
    /// Total number of detections performed
    pub detection_count: u64,
    /// Average inference time in milliseconds
    pub avg_inference_time_ms: f64,
}

/// Request types for detector IPC
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum DetectorRequest {
    /// Initialize the detector with config
    Initialize { config: DetectorConfig },
    /// Detect objects in an image
    Detect { image_data: String },
    /// Get list of supported classes
    GetClasses,
    /// Get detector status
    GetStatus,
}

/// Response types for detector IPC
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum DetectorResponse {
    /// Initialization result
    Initialized { success: bool },
    /// Detection result
    Detected { result: DetectionResult },
    /// List of classes
    Classes { classes: Vec<String> },
    /// Detector status
    Status { status: DetectorStatus },
    /// Error response
    Error { message: String },
}

/// COCO dataset class names (80 classes)
pub const COCO_CLASSES: &[&str] = &[
    "person",
    "bicycle",
    "car",
    "motorcycle",
    "airplane",
    "bus",
    "train",
    "truck",
    "boat",
    "traffic light",
    "fire hydrant",
    "stop sign",
    "parking meter",
    "bench",
    "bird",
    "cat",
    "dog",
    "horse",
    "sheep",
    "cow",
    "elephant",
    "bear",
    "zebra",
    "giraffe",
    "backpack",
    "umbrella",
    "handbag",
    "tie",
    "suitcase",
    "frisbee",
    "skis",
    "snowboard",
    "sports ball",
    "kite",
    "baseball bat",
    "baseball glove",
    "skateboard",
    "surfboard",
    "tennis racket",
    "bottle",
    "wine glass",
    "cup",
    "fork",
    "knife",
    "spoon",
    "bowl",
    "banana",
    "apple",
    "sandwich",
    "orange",
    "broccoli",
    "carrot",
    "hot dog",
    "pizza",
    "donut",
    "cake",
    "chair",
    "couch",
    "potted plant",
    "bed",
    "dining table",
    "toilet",
    "tv",
    "laptop",
    "mouse",
    "remote",
    "keyboard",
    "cell phone",
    "microwave",
    "oven",
    "toaster",
    "sink",
    "refrigerator",
    "book",
    "clock",
    "vase",
    "scissors",
    "teddy bear",
    "hair drier",
    "toothbrush",
];

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_coco_classes_count() {
        assert_eq!(COCO_CLASSES.len(), 80);
    }

    #[test]
    fn test_coco_classes_first_is_person() {
        assert_eq!(COCO_CLASSES[0], "person");
    }

    #[test]
    fn test_coco_classes_last_is_toothbrush() {
        assert_eq!(COCO_CLASSES[79], "toothbrush");
    }

    #[test]
    fn test_coco_classes_contains_common_objects() {
        assert!(COCO_CLASSES.contains(&"car"));
        assert!(COCO_CLASSES.contains(&"dog"));
        assert!(COCO_CLASSES.contains(&"cat"));
    }

    #[test]
    fn test_detector_model_serialization() {
        let model = DetectorModel::YOLOv8_N;
        let json = serde_json::to_string(&model).expect("failed to serialize");
        assert_eq!(json, "\"yolov8_n\"");
    }

    #[test]
    fn test_detector_model_deserialization() {
        let json = "\"yolov8_s\"";
        let model: DetectorModel = serde_json::from_str(json).expect("failed to deserialize");
        assert_eq!(model, DetectorModel::YOLOv8_S);
    }

    #[test]
    fn test_detector_model_display_name() {
        assert_eq!(DetectorModel::YOLOv8_N.display_name(), "YOLOv8-Nano");
        assert_eq!(DetectorModel::YOLOv8_S.display_name(), "YOLOv8-Small");
        assert_eq!(DetectorModel::YOLOv8_M.display_name(), "YOLOv8-Medium");
        assert_eq!(DetectorModel::YOLOv8_L.display_name(), "YOLOv8-Large");
        assert_eq!(DetectorModel::YOLOv8_X.display_name(), "YOLOv8-XLarge");
    }

    #[test]
    fn test_detector_model_param_count() {
        assert!(DetectorModel::YOLOv8_N.param_count() < DetectorModel::YOLOv8_S.param_count());
        assert!(DetectorModel::YOLOv8_S.param_count() < DetectorModel::YOLOv8_M.param_count());
        assert!(DetectorModel::YOLOv8_M.param_count() < DetectorModel::YOLOv8_L.param_count());
        assert!(DetectorModel::YOLOv8_L.param_count() < DetectorModel::YOLOv8_X.param_count());
    }

    #[test]
    fn test_bounding_box_serialization() {
        let bbox = BoundingBox {
            x: 0.25,
            y: 0.30,
            width: 0.50,
            height: 0.40,
        };
        let json = serde_json::to_string(&bbox).expect("failed to serialize");
        assert!(json.contains("0.25"));
        assert!(json.contains("0.3"));
    }

    #[test]
    fn test_bounding_box_normalized_values() {
        let bbox = BoundingBox {
            x: 0.0,
            y: 0.0,
            width: 1.0,
            height: 1.0,
        };
        assert!(bbox.x >= 0.0 && bbox.x <= 1.0);
        assert!(bbox.y >= 0.0 && bbox.y <= 1.0);
        assert!(bbox.width >= 0.0 && bbox.width <= 1.0);
        assert!(bbox.height >= 0.0 && bbox.height <= 1.0);
    }

    #[test]
    fn test_bounding_box_area() {
        let bbox = BoundingBox {
            x: 0.2,
            y: 0.2,
            width: 0.4,
            height: 0.3,
        };
        let area = bbox.area();
        assert!((area - 0.12).abs() < 1e-6);
    }

    #[test]
    fn test_bounding_box_center() {
        let bbox = BoundingBox {
            x: 0.2,
            y: 0.3,
            width: 0.4,
            height: 0.6,
        };
        let (cx, cy) = bbox.center();
        assert_eq!(cx, 0.4);
        assert_eq!(cy, 0.6);
    }

    #[test]
    fn test_bounding_box_intersection() {
        let bbox1 = BoundingBox {
            x: 0.1,
            y: 0.1,
            width: 0.5,
            height: 0.5,
        };
        let bbox2 = BoundingBox {
            x: 0.3,
            y: 0.3,
            width: 0.5,
            height: 0.5,
        };
        let intersection = bbox1.intersection(&bbox2);
        assert_eq!(intersection, 0.09);
    }

    #[test]
    fn test_bounding_box_iou() {
        let bbox1 = BoundingBox {
            x: 0.0,
            y: 0.0,
            width: 0.5,
            height: 0.5,
        };
        let bbox2 = BoundingBox {
            x: 0.0,
            y: 0.0,
            width: 0.5,
            height: 0.5,
        };
        let iou = bbox1.iou(&bbox2);
        assert_eq!(iou, 1.0);
    }

    #[test]
    fn test_bounding_box_no_overlap_iou() {
        let bbox1 = BoundingBox {
            x: 0.0,
            y: 0.0,
            width: 0.2,
            height: 0.2,
        };
        let bbox2 = BoundingBox {
            x: 0.5,
            y: 0.5,
            width: 0.2,
            height: 0.2,
        };
        let iou = bbox1.iou(&bbox2);
        assert_eq!(iou, 0.0);
    }

    #[test]
    fn test_detection_serialization() {
        let detection = Detection {
            bbox: BoundingBox {
                x: 0.1,
                y: 0.2,
                width: 0.3,
                height: 0.4,
            },
            class_id: 0,
            class_name: "person".to_string(),
            confidence: 0.95,
            color: "#FF5733".to_string(),
        };
        let json = serde_json::to_string(&detection).expect("failed to serialize");
        assert!(json.contains("person"));
        assert!(json.contains("0.95"));
    }

    #[test]
    fn test_detection_result_serialization() {
        let result = DetectionResult {
            detections: vec![],
            inference_time_ms: 42.5,
            image_width: 640,
            image_height: 480,
        };
        let json = serde_json::to_string(&result).expect("failed to serialize");
        assert!(json.contains("42.5"));
        assert!(json.contains("640"));
    }

    #[test]
    fn test_detector_config_default() {
        let config = DetectorConfig::default();
        assert_eq!(config.model, DetectorModel::YOLOv8_N);
        assert_eq!(config.confidence_threshold, 0.25);
        assert_eq!(config.nms_threshold, 0.45);
        assert_eq!(config.max_detections, 100);
        assert!(config.use_gpu);
    }

    #[test]
    fn test_detector_config_custom() {
        let config = DetectorConfig {
            model: DetectorModel::YOLOv8_L,
            confidence_threshold: 0.5,
            nms_threshold: 0.5,
            max_detections: 50,
            use_gpu: false,
        };
        assert_eq!(config.model, DetectorModel::YOLOv8_L);
        assert_eq!(config.confidence_threshold, 0.5);
    }

    #[test]
    fn test_detector_config_serialization() {
        let config = DetectorConfig::default();
        let json = serde_json::to_string(&config).expect("failed to serialize");
        let deserialized: DetectorConfig =
            serde_json::from_str(&json).expect("failed to deserialize");
        assert_eq!(config.model, deserialized.model);
        assert_eq!(
            config.confidence_threshold,
            deserialized.confidence_threshold
        );
    }

    #[test]
    fn test_detector_request_initialize() {
        let request = DetectorRequest::Initialize {
            config: DetectorConfig::default(),
        };
        let json = serde_json::to_string(&request).expect("failed to serialize");
        assert!(json.contains("initialize"));
    }

    #[test]
    fn test_detector_request_detect() {
        let request = DetectorRequest::Detect {
            image_data: "base64data".to_string(),
        };
        let json = serde_json::to_string(&request).expect("failed to serialize");
        assert!(json.contains("detect"));
        assert!(json.contains("base64data"));
    }

    #[test]
    fn test_detector_request_get_classes() {
        let request = DetectorRequest::GetClasses;
        let json = serde_json::to_string(&request).expect("failed to serialize");
        assert!(json.contains("get_classes"));
    }

    #[test]
    fn test_detector_request_get_status() {
        let request = DetectorRequest::GetStatus;
        let json = serde_json::to_string(&request).expect("failed to serialize");
        assert!(json.contains("get_status"));
    }

    #[test]
    fn test_detector_response_initialized() {
        let response = DetectorResponse::Initialized { success: true };
        let json = serde_json::to_string(&response).expect("failed to serialize");
        assert!(json.contains("initialized"));
        assert!(json.contains("true"));
    }

    #[test]
    fn test_detector_response_detected() {
        let result = DetectionResult {
            detections: vec![],
            inference_time_ms: 50.0,
            image_width: 800,
            image_height: 600,
        };
        let response = DetectorResponse::Detected { result };
        let json = serde_json::to_string(&response).expect("failed to serialize");
        assert!(json.contains("detected"));
    }

    #[test]
    fn test_detector_response_classes() {
        let classes = vec!["person".to_string(), "car".to_string()];
        let response = DetectorResponse::Classes { classes };
        let json = serde_json::to_string(&response).expect("failed to serialize");
        assert!(json.contains("classes"));
        assert!(json.contains("person"));
    }

    #[test]
    fn test_detector_response_status() {
        let status = DetectorStatus {
            initialized: true,
            model: Some("YOLOv8-Nano".to_string()),
            device: "CPU".to_string(),
            detection_count: 42,
            avg_inference_time_ms: 35.5,
        };
        let response = DetectorResponse::Status { status };
        let json = serde_json::to_string(&response).expect("failed to serialize");
        assert!(json.contains("status"));
        assert!(json.contains("42"));
    }

    #[test]
    fn test_detector_response_error() {
        let response = DetectorResponse::Error {
            message: "Test error".to_string(),
        };
        let json = serde_json::to_string(&response).expect("failed to serialize");
        assert!(json.contains("error"));
        assert!(json.contains("Test error"));
    }

    #[test]
    fn test_detector_status_uninitialized() {
        let status = DetectorStatus {
            initialized: false,
            model: None,
            device: "CPU".to_string(),
            detection_count: 0,
            avg_inference_time_ms: 0.0,
        };
        assert!(!status.initialized);
        assert!(status.model.is_none());
    }
}