128 lines
4.4 KiB
Rust
128 lines
4.4 KiB
Rust
// RTX Segmentation Demo - Mock segmentation backend
|
|
//
|
|
// This crate provides mock implementations for semantic segmentation inference,
|
|
// including realistic region-based mask generation and colormap visualization.
|
|
|
|
mod colormap;
|
|
mod error;
|
|
mod mock_segmentation;
|
|
|
|
pub use colormap::{apply_colormap_with_transparency, blend_with_image, mask_to_rgb};
|
|
pub use error::{Result, SegmentationError};
|
|
pub use mock_segmentation::{
|
|
generate_for_model, generate_mock_segmentation, generate_with_distribution,
|
|
};
|
|
|
|
// Re-export shared types for convenience
|
|
pub use segmentation_shared::{
|
|
ADE20K_CLASSES, CITYSCAPES_CLASSES, ClassInfo, PASCAL_VOC_CLASSES, SegmentRequest,
|
|
SegmentResponse, SegmentationConfig, SegmentationMask, SegmentationModel, SegmentationResult,
|
|
};
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_library_exports_error_types() {
|
|
let _err: SegmentationError = SegmentationError::InvalidDimensions("test".to_string());
|
|
let _result: Result<()> = Ok(());
|
|
}
|
|
|
|
#[test]
|
|
fn test_library_exports_colormap_functions() {
|
|
let mask = SegmentationMask::new(1, 1, vec![0]);
|
|
let classes = &[ClassInfo::new(0, "bg", "#000000")];
|
|
|
|
let _rgb = mask_to_rgb(&mask, classes).unwrap();
|
|
let _rgba = apply_colormap_with_transparency(&mask, classes, true).unwrap();
|
|
|
|
let original = vec![100, 100, 100];
|
|
let _blended = blend_with_image(&original, &mask, classes, 0.5).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn test_library_exports_mock_generation() {
|
|
let config =
|
|
SegmentationConfig::new(SegmentationModel::DeepLabV3, "cpu".to_string(), 21, 512);
|
|
let _result = generate_mock_segmentation(&config, 10, 10, Some(42)).unwrap();
|
|
let _result2 = generate_for_model(SegmentationModel::UNet, 10, 10, Some(42)).unwrap();
|
|
let _result3 = generate_with_distribution(10, 10, &[1.0, 1.0], Some(42)).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn test_library_exports_shared_types() {
|
|
let _model = SegmentationModel::DeepLabV3;
|
|
let _config = SegmentationConfig::new(_model, "cpu".to_string(), 21, 512);
|
|
let _mask = SegmentationMask::new(1, 1, vec![0]);
|
|
let _result = SegmentationResult::new(_mask.clone(), 100.0);
|
|
let _class = ClassInfo::new(0, "bg", "#000000");
|
|
}
|
|
|
|
#[test]
|
|
fn test_library_exports_ipc_types() {
|
|
let _request = SegmentRequest::GetModels;
|
|
let _response = SegmentResponse::error("test".to_string());
|
|
}
|
|
|
|
#[test]
|
|
fn test_library_exports_class_sets() {
|
|
assert!(!PASCAL_VOC_CLASSES.is_empty());
|
|
assert!(!ADE20K_CLASSES.is_empty());
|
|
assert!(!CITYSCAPES_CLASSES.is_empty());
|
|
}
|
|
|
|
#[test]
|
|
fn test_end_to_end_workflow() {
|
|
let config =
|
|
SegmentationConfig::new(SegmentationModel::DeepLabV3, "cpu".to_string(), 21, 512);
|
|
|
|
let result = generate_mock_segmentation(&config, 50, 50, Some(42)).unwrap();
|
|
|
|
assert!(result.validate().is_ok());
|
|
|
|
let rgb = mask_to_rgb(&result.mask, PASCAL_VOC_CLASSES).unwrap();
|
|
assert_eq!(rgb.len(), 50 * 50 * 3);
|
|
|
|
let rgba =
|
|
apply_colormap_with_transparency(&result.mask, PASCAL_VOC_CLASSES, true).unwrap();
|
|
assert_eq!(rgba.len(), 50 * 50 * 4);
|
|
}
|
|
|
|
#[test]
|
|
fn test_end_to_end_with_blending() {
|
|
let config = SegmentationConfig::new(SegmentationModel::UNet, "cpu".to_string(), 21, 512);
|
|
|
|
let result = generate_mock_segmentation(&config, 20, 20, Some(99)).unwrap();
|
|
|
|
let original_image = vec![128u8; 20 * 20 * 3];
|
|
|
|
let blended =
|
|
blend_with_image(&original_image, &result.mask, PASCAL_VOC_CLASSES, 0.6).unwrap();
|
|
assert_eq!(blended.len(), original_image.len());
|
|
|
|
// All values are valid u8 by type (no need to check bounds)
|
|
}
|
|
|
|
#[test]
|
|
fn test_multiple_models_generate_different_class_counts() {
|
|
let models = [
|
|
(SegmentationModel::DeepLabV3, 21),
|
|
(SegmentationModel::SegFormer, 150),
|
|
(SegmentationModel::UNet, 21),
|
|
(SegmentationModel::FCN, 21),
|
|
];
|
|
|
|
for (model, expected_classes) in models {
|
|
let result = generate_for_model(model, 30, 30, Some(42)).unwrap();
|
|
assert_eq!(
|
|
result.class_confidences.len(),
|
|
expected_classes,
|
|
"model {:?} should have {} classes",
|
|
model,
|
|
expected_classes
|
|
);
|
|
}
|
|
}
|
|
}
|