424 lines
13 KiB
Rust
424 lines
13 KiB
Rust
//! Mock object detector for demonstration purposes
|
|
//!
|
|
//! Generates realistic-looking detections without requiring a real model.
|
|
|
|
use object_detector_shared::{
|
|
BoundingBox, COCO_CLASSES, Detection, DetectionResult, DetectorConfig,
|
|
};
|
|
use rand::{Rng, SeedableRng};
|
|
|
|
use crate::error::{DetectorError, Result};
|
|
|
|
/// Class colors for visualization (hex strings)
|
|
const CLASS_COLORS: &[&str] = &[
|
|
"#FF6B6B", "#4ECDC4", "#45B7D1", "#FFA07A", "#98D8C8", "#F7DC6F", "#BB8FCE", "#85C1E2",
|
|
"#F8B739", "#52B788", "#E63946", "#06FFA5", "#1D3557", "#F77F00", "#06D6A0", "#118AB2",
|
|
"#EF476F", "#FFD166", "#06A77D", "#073B4C",
|
|
];
|
|
|
|
/// Get a color for a class ID
|
|
fn get_class_color(class_id: usize) -> String {
|
|
CLASS_COLORS[class_id % CLASS_COLORS.len()].to_string()
|
|
}
|
|
|
|
/// Generate mock detections for an image
|
|
pub fn generate_mock_detections(
|
|
image_width: usize,
|
|
image_height: usize,
|
|
config: &DetectorConfig,
|
|
seed: Option<u64>,
|
|
) -> Result<DetectionResult> {
|
|
if image_width == 0 || image_height == 0 {
|
|
return Err(DetectorError::InvalidImage(
|
|
"Image dimensions cannot be zero".to_string(),
|
|
));
|
|
}
|
|
|
|
let mut rng = if let Some(seed) = seed {
|
|
rand::rngs::StdRng::seed_from_u64(seed)
|
|
} else {
|
|
rand::rngs::StdRng::from_entropy()
|
|
};
|
|
|
|
let start_time = std::time::Instant::now();
|
|
|
|
// Generate 3-10 random detections
|
|
let num_detections = rng.gen_range(3..=10);
|
|
let mut detections = Vec::new();
|
|
|
|
// Common classes with higher probability
|
|
let common_classes = [0, 2, 15, 16]; // person, car, cat, dog
|
|
|
|
for _ in 0..num_detections {
|
|
// Use common classes 70% of the time
|
|
let class_id = if rng.r#gen::<f32>() < 0.7 {
|
|
common_classes[rng.gen_range(0..common_classes.len())]
|
|
} else {
|
|
rng.gen_range(0..COCO_CLASSES.len())
|
|
};
|
|
|
|
// Generate random bbox with reasonable size
|
|
let width = rng.gen_range(0.1..0.4);
|
|
let height = rng.gen_range(0.1..0.4);
|
|
let x = rng.gen_range(0.0..(1.0 - width));
|
|
let y = rng.gen_range(0.0..(1.0 - height));
|
|
|
|
let bbox = BoundingBox {
|
|
x,
|
|
y,
|
|
width,
|
|
height,
|
|
};
|
|
|
|
// Generate realistic confidence (higher for common classes)
|
|
let confidence = if common_classes.contains(&class_id) {
|
|
rng.gen_range(0.6..0.99)
|
|
} else {
|
|
rng.gen_range(0.3..0.8)
|
|
};
|
|
|
|
// Only add if above confidence threshold
|
|
if confidence >= config.confidence_threshold {
|
|
detections.push(Detection {
|
|
bbox,
|
|
class_id,
|
|
class_name: COCO_CLASSES[class_id].to_string(),
|
|
confidence,
|
|
color: get_class_color(class_id),
|
|
});
|
|
}
|
|
}
|
|
|
|
// Sort by confidence descending
|
|
detections.sort_by(|a, b| b.confidence.partial_cmp(&a.confidence).unwrap());
|
|
|
|
// Apply NMS
|
|
detections = apply_nms(detections, config.nms_threshold);
|
|
|
|
// Limit to max_detections
|
|
detections.truncate(config.max_detections);
|
|
|
|
let inference_time_ms = start_time.elapsed().as_secs_f64() * 1000.0;
|
|
|
|
Ok(DetectionResult {
|
|
detections,
|
|
inference_time_ms,
|
|
image_width,
|
|
image_height,
|
|
})
|
|
}
|
|
|
|
/// Apply Non-Maximum Suppression
|
|
pub fn apply_nms(mut detections: Vec<Detection>, threshold: f32) -> Vec<Detection> {
|
|
let mut result = Vec::new();
|
|
|
|
while !detections.is_empty() {
|
|
// Take the detection with highest confidence
|
|
let best = detections.remove(0);
|
|
|
|
// Keep detections that don't overlap too much with the best one
|
|
detections.retain(|det| {
|
|
// Keep if different class OR low overlap
|
|
det.class_id != best.class_id || det.bbox.iou(&best.bbox) < threshold
|
|
});
|
|
|
|
result.push(best);
|
|
}
|
|
|
|
result
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use approx::assert_relative_eq;
|
|
|
|
#[test]
|
|
fn test_get_class_color_returns_valid_hex() {
|
|
let color = get_class_color(0);
|
|
assert!(color.starts_with('#'));
|
|
assert_eq!(color.len(), 7);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_class_color_wraps_around() {
|
|
let color1 = get_class_color(0);
|
|
let color2 = get_class_color(CLASS_COLORS.len());
|
|
assert_eq!(color1, color2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_zero_width_fails() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(0, 480, &config, None);
|
|
assert!(result.is_err());
|
|
assert!(
|
|
result
|
|
.unwrap_err()
|
|
.to_string()
|
|
.contains("dimensions cannot be zero")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_zero_height_fails() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(640, 0, &config, None);
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_with_seed_is_deterministic() {
|
|
let config = DetectorConfig::default();
|
|
let result1 = generate_mock_detections(640, 480, &config, Some(42)).unwrap();
|
|
let result2 = generate_mock_detections(640, 480, &config, Some(42)).unwrap();
|
|
|
|
assert_eq!(result1.detections.len(), result2.detections.len());
|
|
assert_eq!(result1.image_width, result2.image_width);
|
|
assert_eq!(result1.image_height, result2.image_height);
|
|
|
|
for (d1, d2) in result1.detections.iter().zip(result2.detections.iter()) {
|
|
assert_eq!(d1.class_id, d2.class_id);
|
|
assert_relative_eq!(d1.confidence, d2.confidence);
|
|
assert_relative_eq!(d1.bbox.x, d2.bbox.x);
|
|
assert_relative_eq!(d1.bbox.y, d2.bbox.y);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_returns_valid_result() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(640, 480, &config, Some(123)).unwrap();
|
|
|
|
assert!(result.detections.len() >= 0);
|
|
assert!(result.detections.len() <= config.max_detections);
|
|
assert_eq!(result.image_width, 640);
|
|
assert_eq!(result.image_height, 480);
|
|
assert!(result.inference_time_ms >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_respects_confidence_threshold() {
|
|
let mut config = DetectorConfig::default();
|
|
config.confidence_threshold = 0.9;
|
|
let result = generate_mock_detections(640, 480, &config, Some(456)).unwrap();
|
|
|
|
for detection in &result.detections {
|
|
assert!(detection.confidence >= config.confidence_threshold);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_respects_max_detections() {
|
|
let mut config = DetectorConfig::default();
|
|
config.max_detections = 3;
|
|
config.confidence_threshold = 0.1; // Low threshold to get more detections
|
|
let result = generate_mock_detections(640, 480, &config, Some(789)).unwrap();
|
|
|
|
assert!(result.detections.len() <= config.max_detections);
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_bbox_in_valid_range() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(640, 480, &config, Some(101)).unwrap();
|
|
|
|
for detection in &result.detections {
|
|
assert!(detection.bbox.x >= 0.0 && detection.bbox.x <= 1.0);
|
|
assert!(detection.bbox.y >= 0.0 && detection.bbox.y <= 1.0);
|
|
assert!(detection.bbox.width >= 0.0 && detection.bbox.width <= 1.0);
|
|
assert!(detection.bbox.height >= 0.0 && detection.bbox.height <= 1.0);
|
|
assert!(detection.bbox.x + detection.bbox.width <= 1.0);
|
|
assert!(detection.bbox.y + detection.bbox.height <= 1.0);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_class_ids_valid() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(640, 480, &config, Some(202)).unwrap();
|
|
|
|
for detection in &result.detections {
|
|
assert!(detection.class_id < COCO_CLASSES.len());
|
|
assert_eq!(detection.class_name, COCO_CLASSES[detection.class_id]);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_has_colors() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(640, 480, &config, Some(303)).unwrap();
|
|
|
|
for detection in &result.detections {
|
|
assert!(detection.color.starts_with('#'));
|
|
assert_eq!(detection.color.len(), 7);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_generate_mock_detections_sorted_by_confidence() {
|
|
let config = DetectorConfig::default();
|
|
let result = generate_mock_detections(640, 480, &config, Some(404)).unwrap();
|
|
|
|
for i in 1..result.detections.len() {
|
|
assert!(result.detections[i - 1].confidence >= result.detections[i].confidence);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_nms_empty_list() {
|
|
let detections = Vec::new();
|
|
let result = apply_nms(detections, 0.5);
|
|
assert!(result.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_nms_single_detection() {
|
|
let detection = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.1,
|
|
y: 0.1,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.9,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
let result = apply_nms(vec![detection.clone()], 0.5);
|
|
assert_eq!(result.len(), 1);
|
|
assert_eq!(result[0].class_id, detection.class_id);
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_nms_removes_overlapping_same_class() {
|
|
let det1 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.1,
|
|
y: 0.1,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.9,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
let det2 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.15,
|
|
y: 0.15,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.8,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
|
|
let result = apply_nms(vec![det1, det2], 0.3);
|
|
assert_eq!(result.len(), 1);
|
|
assert_eq!(result[0].confidence, 0.9);
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_nms_keeps_different_classes() {
|
|
let det1 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.1,
|
|
y: 0.1,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.9,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
let det2 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.15,
|
|
y: 0.15,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 1,
|
|
class_name: "bicycle".to_string(),
|
|
confidence: 0.8,
|
|
color: "#00FF00".to_string(),
|
|
};
|
|
|
|
let result = apply_nms(vec![det1, det2], 0.3);
|
|
assert_eq!(result.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_nms_keeps_non_overlapping_same_class() {
|
|
let det1 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.1,
|
|
y: 0.1,
|
|
width: 0.1,
|
|
height: 0.1,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.9,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
let det2 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.5,
|
|
y: 0.5,
|
|
width: 0.1,
|
|
height: 0.1,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.8,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
|
|
let result = apply_nms(vec![det1, det2], 0.5);
|
|
assert_eq!(result.len(), 2);
|
|
}
|
|
|
|
#[test]
|
|
fn test_apply_nms_respects_threshold() {
|
|
let det1 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.1,
|
|
y: 0.1,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.9,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
let det2 = Detection {
|
|
bbox: BoundingBox {
|
|
x: 0.12,
|
|
y: 0.12,
|
|
width: 0.2,
|
|
height: 0.2,
|
|
},
|
|
class_id: 0,
|
|
class_name: "person".to_string(),
|
|
confidence: 0.8,
|
|
color: "#FF0000".to_string(),
|
|
};
|
|
|
|
// High threshold - should keep both
|
|
let result = apply_nms(vec![det1.clone(), det2.clone()], 0.9);
|
|
assert_eq!(result.len(), 2);
|
|
|
|
// Low threshold - should remove overlapping
|
|
let result = apply_nms(vec![det1, det2], 0.1);
|
|
assert_eq!(result.len(), 1);
|
|
}
|
|
}
|