449 lines
12 KiB
Rust
449 lines
12 KiB
Rust
//! # RTX Vision Advanced
|
|
//!
|
|
//! State-of-the-art computer vision library providing production-ready implementations of:
|
|
//! - Object detection (YOLO v8/v9, R-CNN family)
|
|
//! - Instance and panoptic segmentation
|
|
//! - Medical imaging with DICOM support
|
|
//! - Autonomous vehicle perception
|
|
//! - Real-time optimization and deployment
|
|
//!
|
|
//! ## Features
|
|
//!
|
|
//! - **Detection**: YOLO v8/v9, Faster R-CNN, Mask R-CNN with GPU acceleration
|
|
//! - **Medical**: DICOM processing, organ segmentation, radiological AI
|
|
//! - **Autonomous**: LiDAR processing, sensor fusion, path planning integration
|
|
//! - **Production**: Real-time inference, model quantization, edge deployment
|
|
//! - **Advanced**: Few-shot detection, vision-language models, self-supervised learning
|
|
|
|
use anyhow::Result;
|
|
use tracing::info;
|
|
|
|
pub mod autonomous;
|
|
pub mod detection;
|
|
pub mod medical;
|
|
pub mod segmentation;
|
|
// pub mod tracking; // Module to be implemented
|
|
pub mod error;
|
|
pub mod production;
|
|
pub mod tensor_utils;
|
|
pub mod utils;
|
|
|
|
pub use error::{VisionError, VisionResult};
|
|
|
|
// Re-export core types for convenience
|
|
pub use rtx_autograd::{Variable, backward};
|
|
pub use rtx_tensor::{DType, Device, Tensor};
|
|
|
|
/// Configuration for vision models
|
|
#[derive(Debug, Clone)]
|
|
pub struct VisionConfig {
|
|
/// Target device for computation
|
|
pub device: Device,
|
|
/// Data type for computation
|
|
pub dtype: DType,
|
|
/// Whether to use mixed precision
|
|
pub mixed_precision: bool,
|
|
/// Batch size for inference
|
|
pub batch_size: usize,
|
|
/// Input image size (height, width)
|
|
pub input_size: (usize, usize),
|
|
/// Number of classes
|
|
pub num_classes: usize,
|
|
/// Confidence threshold for detection
|
|
pub confidence_threshold: f32,
|
|
/// NMS threshold for detection
|
|
pub nms_threshold: f32,
|
|
}
|
|
|
|
impl Default for VisionConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
device: Device::default(),
|
|
dtype: DType::F32,
|
|
mixed_precision: false,
|
|
batch_size: 1,
|
|
input_size: (640, 640),
|
|
num_classes: 80, // COCO classes
|
|
confidence_threshold: 0.25,
|
|
nms_threshold: 0.45,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Model weights information
|
|
#[derive(Debug, Clone)]
|
|
pub struct ModelWeights {
|
|
/// Model name
|
|
pub name: String,
|
|
/// Model version
|
|
pub version: String,
|
|
/// URL or path to weights
|
|
pub weights_path: String,
|
|
/// Model configuration
|
|
pub config: serde_json::Value,
|
|
/// Model input size
|
|
pub input_size: (usize, usize),
|
|
/// Number of classes
|
|
pub num_classes: usize,
|
|
}
|
|
|
|
/// Performance metrics for evaluation
|
|
#[derive(Debug, Clone)]
|
|
pub struct PerformanceMetrics {
|
|
/// Mean Average Precision (mAP)
|
|
pub map: f32,
|
|
/// mAP at IoU 0.5
|
|
pub map50: f32,
|
|
/// mAP at IoU 0.75
|
|
pub map75: f32,
|
|
/// Average precision per class
|
|
pub class_ap: Vec<f32>,
|
|
/// Inference time in milliseconds
|
|
pub inference_time_ms: f32,
|
|
/// Memory usage in MB
|
|
pub memory_usage_mb: f32,
|
|
/// Frames per second
|
|
pub fps: f32,
|
|
}
|
|
|
|
impl Default for PerformanceMetrics {
|
|
fn default() -> Self {
|
|
Self {
|
|
map: 0.0,
|
|
map50: 0.0,
|
|
map75: 0.0,
|
|
class_ap: Vec::new(),
|
|
inference_time_ms: 0.0,
|
|
memory_usage_mb: 0.0,
|
|
fps: 0.0,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Bounding box representation
|
|
#[derive(Debug, Clone, PartialEq)]
|
|
pub struct BoundingBox {
|
|
/// X coordinate of top-left corner
|
|
pub x: f32,
|
|
/// Y coordinate of top-left corner
|
|
pub y: f32,
|
|
/// Width of the box
|
|
pub width: f32,
|
|
/// Height of the box
|
|
pub height: f32,
|
|
/// Confidence score
|
|
pub confidence: f32,
|
|
/// Class ID
|
|
pub class_id: usize,
|
|
/// Class name (optional)
|
|
pub class_name: Option<String>,
|
|
}
|
|
|
|
impl BoundingBox {
|
|
/// Create a new bounding box
|
|
pub fn new(x: f32, y: f32, width: f32, height: f32, confidence: f32, class_id: usize) -> Self {
|
|
Self {
|
|
x,
|
|
y,
|
|
width,
|
|
height,
|
|
confidence,
|
|
class_id,
|
|
class_name: None,
|
|
}
|
|
}
|
|
|
|
/// Calculate the area of the bounding box
|
|
pub fn area(&self) -> f32 {
|
|
self.width * self.height
|
|
}
|
|
|
|
/// Calculate IoU with another bounding box
|
|
pub fn iou(&self, other: &Self) -> f32 {
|
|
let x1_max = self.x.max(other.x);
|
|
let y1_max = self.y.max(other.y);
|
|
let x2_min = (self.x + self.width).min(other.x + other.width);
|
|
let y2_min = (self.y + self.height).min(other.y + other.height);
|
|
|
|
if x2_min <= x1_max || y2_min <= y1_max {
|
|
return 0.0;
|
|
}
|
|
|
|
let intersection = (x2_min - x1_max) * (y2_min - y1_max);
|
|
let union = self.area() + other.area() - intersection;
|
|
|
|
intersection / union
|
|
}
|
|
|
|
/// Convert to center format (cx, cy, width, height)
|
|
pub fn to_center_format(&self) -> (f32, f32, f32, f32) {
|
|
let cx = self.x + self.width / 2.0;
|
|
let cy = self.y + self.height / 2.0;
|
|
(cx, cy, self.width, self.height)
|
|
}
|
|
}
|
|
|
|
/// Detection result containing bounding boxes and metadata
|
|
#[derive(Debug, Clone)]
|
|
pub struct DetectionResult {
|
|
/// List of detected bounding boxes
|
|
pub boxes: Vec<BoundingBox>,
|
|
/// Input image dimensions
|
|
pub image_size: (usize, usize),
|
|
/// Processing time in milliseconds
|
|
pub processing_time_ms: f32,
|
|
/// Model name used for detection
|
|
pub model_name: String,
|
|
}
|
|
|
|
impl DetectionResult {
|
|
/// Create a new detection result
|
|
pub fn new(
|
|
boxes: Vec<BoundingBox>,
|
|
image_size: (usize, usize),
|
|
processing_time_ms: f32,
|
|
model_name: String,
|
|
) -> Self {
|
|
Self {
|
|
boxes,
|
|
image_size,
|
|
processing_time_ms,
|
|
model_name,
|
|
}
|
|
}
|
|
|
|
/// Filter detections by confidence threshold
|
|
pub fn filter_by_confidence(&mut self, threshold: f32) {
|
|
self.boxes.retain(|bbox| bbox.confidence >= threshold);
|
|
}
|
|
|
|
/// Get detections for a specific class
|
|
pub fn get_class_detections(&self, class_id: usize) -> Vec<&BoundingBox> {
|
|
self.boxes
|
|
.iter()
|
|
.filter(|bbox| bbox.class_id == class_id)
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
/// Segmentation mask representation
|
|
#[derive(Debug, Clone)]
|
|
pub struct SegmentationMask {
|
|
/// Mask data (height x width)
|
|
pub mask: Tensor,
|
|
/// Class ID for the mask
|
|
pub class_id: usize,
|
|
/// Confidence score
|
|
pub confidence: f32,
|
|
/// Bounding box around the mask
|
|
pub bbox: Option<BoundingBox>,
|
|
}
|
|
|
|
impl SegmentationMask {
|
|
/// Create a new segmentation mask
|
|
pub fn new(mask: Tensor, class_id: usize, confidence: f32) -> Self {
|
|
Self {
|
|
mask,
|
|
class_id,
|
|
confidence,
|
|
bbox: None,
|
|
}
|
|
}
|
|
|
|
/// Calculate mask area (number of pixels)
|
|
pub fn area(&self) -> Result<usize> {
|
|
let mask_data = self.mask.to_vec()?;
|
|
Ok(mask_data.iter().filter(|&&x| x > 0.0).count())
|
|
}
|
|
|
|
/// Calculate IoU with another mask
|
|
pub fn mask_iou(&self, other: &Self) -> Result<f32> {
|
|
let mask1 = &self.mask;
|
|
let mask2 = &other.mask;
|
|
|
|
// Compute intersection and union using multiplication and addition
|
|
// For binary masks, intersection is element-wise multiplication
|
|
let intersection = mask1.mul(mask2)?;
|
|
let union_tensor = mask1.add(mask2)?;
|
|
// Union = sum(mask1 + mask2 - intersection)
|
|
let union = union_tensor.sub(&intersection)?;
|
|
|
|
use tensor_utils::TensorExt;
|
|
let intersection_sum = intersection.mean_dim(&[], false)?;
|
|
let union_sum = union.mean_dim(&[], false)?;
|
|
|
|
let intersection_val: f32 = intersection_sum.to_scalar()?;
|
|
let union_val: f32 = union_sum.to_scalar()?;
|
|
|
|
if union_val == 0.0 {
|
|
return Ok(0.0);
|
|
}
|
|
|
|
Ok(intersection_val / union_val)
|
|
}
|
|
}
|
|
|
|
/// Segmentation result containing masks and metadata
|
|
#[derive(Debug, Clone)]
|
|
pub struct SegmentationResult {
|
|
/// List of segmentation masks
|
|
pub masks: Vec<SegmentationMask>,
|
|
/// Input image dimensions
|
|
pub image_size: (usize, usize),
|
|
/// Processing time in milliseconds
|
|
pub processing_time_ms: f32,
|
|
/// Model name used for segmentation
|
|
pub model_name: String,
|
|
}
|
|
|
|
impl SegmentationResult {
|
|
/// Create a new segmentation result
|
|
pub fn new(
|
|
masks: Vec<SegmentationMask>,
|
|
image_size: (usize, usize),
|
|
processing_time_ms: f32,
|
|
model_name: String,
|
|
) -> Self {
|
|
Self {
|
|
masks,
|
|
image_size,
|
|
processing_time_ms,
|
|
model_name,
|
|
}
|
|
}
|
|
|
|
/// Get masks for a specific class
|
|
pub fn get_class_masks(&self, class_id: usize) -> Vec<&SegmentationMask> {
|
|
self.masks
|
|
.iter()
|
|
.filter(|mask| mask.class_id == class_id)
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
/// Initialize the vision advanced library
|
|
pub fn init() -> Result<()> {
|
|
info!("Initializing RTX Vision Advanced library");
|
|
|
|
// Initialize logging if not already done
|
|
// if let Err(_) = tracing_subscriber::fmt::try_init() {
|
|
// debug!("Tracing subscriber already initialized");
|
|
// }
|
|
|
|
// Initialize GPU if available
|
|
#[cfg(feature = "cuda")]
|
|
{
|
|
use rtx_tensor::cuda;
|
|
if cuda::is_available() {
|
|
info!("CUDA is available");
|
|
} else {
|
|
warn!("CUDA is not available, falling back to CPU");
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "metal")]
|
|
{
|
|
info!("Metal backend initialized");
|
|
}
|
|
|
|
info!("RTX Vision Advanced library initialized successfully");
|
|
Ok(())
|
|
}
|
|
|
|
/// Get available models for different tasks
|
|
pub fn list_available_models() -> Vec<ModelWeights> {
|
|
vec![
|
|
// YOLO models
|
|
ModelWeights {
|
|
name: "yolov8n".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
weights_path: "models/yolov8n.safetensors".to_string(),
|
|
config: serde_json::json!({
|
|
"architecture": "yolov8",
|
|
"variant": "nano"
|
|
}),
|
|
input_size: (640, 640),
|
|
num_classes: 80,
|
|
},
|
|
ModelWeights {
|
|
name: "yolov9c".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
weights_path: "models/yolov9c.safetensors".to_string(),
|
|
config: serde_json::json!({
|
|
"architecture": "yolov9",
|
|
"variant": "compact"
|
|
}),
|
|
input_size: (640, 640),
|
|
num_classes: 80,
|
|
},
|
|
// R-CNN models
|
|
ModelWeights {
|
|
name: "faster_rcnn_r50_fpn".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
weights_path: "models/faster_rcnn_r50_fpn.safetensors".to_string(),
|
|
config: serde_json::json!({
|
|
"architecture": "faster_rcnn",
|
|
"backbone": "resnet50",
|
|
"neck": "fpn"
|
|
}),
|
|
input_size: (800, 1333),
|
|
num_classes: 91, // COCO + background
|
|
},
|
|
ModelWeights {
|
|
name: "mask_rcnn_r50_fpn".to_string(),
|
|
version: "1.0.0".to_string(),
|
|
weights_path: "models/mask_rcnn_r50_fpn.safetensors".to_string(),
|
|
config: serde_json::json!({
|
|
"architecture": "mask_rcnn",
|
|
"backbone": "resnet50",
|
|
"neck": "fpn"
|
|
}),
|
|
input_size: (800, 1333),
|
|
num_classes: 91,
|
|
},
|
|
]
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn test_bounding_box_iou() {
|
|
let box1 = BoundingBox::new(10.0, 10.0, 20.0, 20.0, 0.9, 0);
|
|
let box2 = BoundingBox::new(15.0, 15.0, 20.0, 20.0, 0.8, 0);
|
|
|
|
let iou = box1.iou(&box2);
|
|
assert!(iou > 0.0 && iou < 1.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_bounding_box_area() {
|
|
let bbox = BoundingBox::new(0.0, 0.0, 10.0, 20.0, 0.9, 0);
|
|
assert_eq!(bbox.area(), 200.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_vision_config_default() {
|
|
let config = VisionConfig::default();
|
|
assert_eq!(config.input_size, (640, 640));
|
|
assert_eq!(config.num_classes, 80);
|
|
assert_eq!(config.confidence_threshold, 0.25);
|
|
assert_eq!(config.nms_threshold, 0.45);
|
|
}
|
|
|
|
#[test]
|
|
fn test_detection_result_filter() {
|
|
let boxes = vec![
|
|
BoundingBox::new(0.0, 0.0, 10.0, 10.0, 0.9, 0),
|
|
BoundingBox::new(10.0, 10.0, 10.0, 10.0, 0.3, 1),
|
|
BoundingBox::new(20.0, 20.0, 10.0, 10.0, 0.7, 2),
|
|
];
|
|
|
|
let mut result = DetectionResult::new(boxes, (640, 640), 10.0, "test_model".to_string());
|
|
|
|
result.filter_by_confidence(0.5);
|
|
assert_eq!(result.boxes.len(), 2);
|
|
}
|
|
}
|