Add SAM3 (Segment Anything 3) interactive segmentation

src-tauri/Cargo.toml

@@ -35,6 +35,7 @@ image = "0.24"
 ndarray = "0.16.1"
 tokenizers = { version = "0.22.1" }
 thiserror = "1.0"
+base64 = "0.22"
 
 # ONNX Runtime with platform-specific execution providers
 [target.'cfg(target_os = "windows")'.dependencies]

src-tauri/src/lib.rs

@@ -11,18 +11,24 @@ pub mod model_manager;
 pub mod inference_engine;
 mod audio_decoder;
 pub mod vlm_onnx;
+pub mod sam3_engine;
+pub mod sam3_image_processor;
 
 use model_manager::{ModelManager, DownloadProgress};
 use inference_engine::{InferenceEngine, SharedInferenceEngine, create_shared_engine};
 use audio_decoder::decode_audio_file;
 use vlm_onnx::VlmOnnx;
+use sam3_engine::{Sam3Engine, Sam3SegmentResult};
 
 // Download cancellation state
 pub type DownloadCancellation = Arc<AtomicBool>;
 
 // Shared ONNX engine
 pub type SharedOnnxEngine = Arc<tokio::sync::Mutex<Option<VlmOnnx>>>;
 
+// Shared SAM3 engine
+pub type SharedSam3Engine = Arc<tokio::sync::Mutex<Option<Sam3Engine>>>;
+
 // Chat message for conversation history
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ChatMessage {
@@ -400,13 +406,159 @@ async fn is_onnx_model_downloaded(model_id: String) -> Result<bool, String> {
         .map_err(|e| e.to_string())
 }
 
+// ============================================================================
+// SAM3 Commands
+// ============================================================================
+
+// Load SAM3 models
+#[tauri::command]
+async fn sam3_load(
+    sam3_engine: State<'_, SharedSam3Engine>,
+) -> Result<(), String> {
+    println!("Loading SAM3 models...");
+
+    let manager = ModelManager::new().map_err(|e| e.to_string())?;
+    let (vision_path, decoder_path) = manager
+        .get_sam3_model_paths()
+        .await
+        .map_err(|e| e.to_string())?;
+
+    println!("SAM3 Vision: {:?}", vision_path);
+    println!("SAM3 Decoder: {:?}", decoder_path);
+
+    // Load the SAM3 model in a blocking task
+    let engine = tokio::task::spawn_blocking(move || {
+        Sam3Engine::load(&vision_path, &decoder_path)
+    })
+    .await
+    .map_err(|e| format!("Task join error: {}", e))?
+    .map_err(|e| e.to_string())?;
+
+    // Store the engine
+    let mut engine_lock = sam3_engine.lock().await;
+    *engine_lock = Some(engine);
+
+    println!("SAM3 loaded successfully");
+    Ok(())
+}
+
+// Unload SAM3 models to free GPU memory
+#[tauri::command]
+async fn sam3_unload(
+    sam3_engine: State<'_, SharedSam3Engine>,
+) -> Result<(), String> {
+    println!("Unloading SAM3...");
+
+    let mut engine_lock = sam3_engine.lock().await;
+    *engine_lock = None;
+
+    println!("SAM3 unloaded");
+    Ok(())
+}
+
+// Check if SAM3 is loaded
+#[tauri::command]
+async fn sam3_is_loaded(
+    sam3_engine: State<'_, SharedSam3Engine>,
+) -> Result<bool, String> {
+    let engine_lock = sam3_engine.lock().await;
+    Ok(engine_lock.is_some())
+}
+
+// Encode an image with SAM3 vision encoder
+#[tauri::command]
+async fn sam3_encode_image(
+    image_data: Vec<u8>,
+    sam3_engine: State<'_, SharedSam3Engine>,
+) -> Result<(), String> {
+    println!("Encoding image with SAM3...");
+
+    let mut engine_lock = sam3_engine.lock().await;
+    let engine = engine_lock.as_mut()
+        .ok_or_else(|| "SAM3 not loaded".to_string())?;
+
+    // Clone the image data for the blocking task
+    let image_data_clone = image_data.clone();
+
+    // We need to work around the borrow checker by taking ownership
+    let mut engine_taken = engine_lock.take()
+        .ok_or_else(|| "SAM3 not loaded".to_string())?;
+    drop(engine_lock);
+
+    // Run encoding in a blocking task
+    let result = tokio::task::spawn_blocking(move || {
+        engine_taken.encode_image(&image_data_clone)?;
+        Ok::<_, anyhow::Error>(engine_taken)
+    })
+    .await
+    .map_err(|e| format!("Task join error: {}", e))?;
+
+    // Put the engine back
+    let engine_back = result.map_err(|e| e.to_string())?;
+    let mut engine_lock = sam3_engine.lock().await;
+    *engine_lock = Some(engine_back);
+
+    println!("Image encoded");
+    Ok(())
+}
+
+// Run SAM3 segmentation with prompts
+#[tauri::command]
+async fn sam3_segment(
+    points: Vec<(f32, f32)>,
+    labels: Vec<i64>,
+    bbox: Option<(f32, f32, f32, f32)>,
+    sam3_engine: State<'_, SharedSam3Engine>,
+) -> Result<Sam3SegmentResult, String> {
+    let mut engine_lock = sam3_engine.lock().await;
+    let engine = engine_lock.as_mut()
+        .ok_or_else(|| "SAM3 not loaded".to_string())?;
+
+    // Run segmentation (fast, can run in async context)
+    engine.segment(&points, &labels, bbox)
+        .map_err(|e| e.to_string())
+}
+
+// Check if SAM3 model is downloaded
+#[tauri::command]
+async fn sam3_is_model_downloaded() -> Result<bool, String> {
+    let manager = ModelManager::new().map_err(|e| e.to_string())?;
+    manager
+        .is_sam3_downloaded()
+        .await
+        .map_err(|e| e.to_string())
+}
+
+// Download SAM3 model
+#[tauri::command]
+async fn sam3_download_model(
+    app: tauri::AppHandle,
+    cancellation: State<'_, DownloadCancellation>,
+) -> Result<(), String> {
+    cancellation.store(false, Ordering::SeqCst);
+
+    let manager = ModelManager::new().map_err(|e| e.to_string())?;
+    let cancel_flag = cancellation.inner().clone();
+
+    manager
+        .download_sam3_model(
+            move |progress: DownloadProgress| {
+                let _ = app.emit("download-progress", &progress);
+            },
+            cancel_flag,
+        )
+        .await
+        .map_err(|e| e.to_string())
+}
+
 #[cfg_attr(mobile, tauri::mobile_entry_point)]
 pub fn run() {
     tauri::Builder::default()
         .plugin(tauri_plugin_opener::init())
         .plugin(tauri_plugin_dialog::init())
         .manage(create_shared_engine())
         .manage(Arc::new(tokio::sync::Mutex::new(None::<VlmOnnx>))) // ONNX engine
+        .manage(Arc::new(tokio::sync::Mutex::new(None::<Sam3Engine>))) // SAM3 engine
         .manage(Arc::new(AtomicBool::new(false))) // Download cancellation flag
         .invoke_handler(tauri::generate_handler![
             greet,
@@ -424,7 +576,15 @@ pub fn run() {
             download_onnx_model,
             load_onnx_model,
             generate_onnx_response,
-            is_onnx_model_downloaded
+            is_onnx_model_downloaded,
+            // SAM3 commands
+            sam3_load,
+            sam3_unload,
+            sam3_is_loaded,
+            sam3_encode_image,
+            sam3_segment,
+            sam3_is_model_downloaded,
+            sam3_download_model
         ])
         .setup(|app| {
             // Create menu

src-tauri/src/model_manager.rs

@@ -620,6 +620,98 @@ impl ModelManager {
             Err(_) => Ok(false),
         }
     }
+
+    // ============================================================================
+    // SAM3 Model Management
+    // ============================================================================
+
+    /// SAM3 model directory name
+    const SAM3_MODEL_ID: &'static str = "sam3-tracker";
+
+    /// Get SAM3 model paths (vision encoder and decoder)
+    pub async fn get_sam3_model_paths(&self) -> Result<(PathBuf, PathBuf)> {
+        let sam3_dir = self.models_dir.join(Self::SAM3_MODEL_ID);
+
+        if !sam3_dir.exists() {
+            return Err(anyhow!("SAM3 model not found. Please download it first."));
+        }
+
+        // Look for ONNX files
+        let mut vision_encoder: Option<PathBuf> = None;
+        let mut decoder: Option<PathBuf> = None;
+
+        let mut entries = tokio::fs::read_dir(&sam3_dir).await?;
+        while let Some(entry) = entries.next_entry().await? {
+            if let Some(filename) = entry.file_name().to_str() {
+                let path = entry.path();
+                let filename_lower = filename.to_lowercase();
+
+                if filename_lower.contains("vision_encoder") && filename_lower.ends_with(".onnx") {
+                    vision_encoder = Some(path);
+                } else if (filename_lower.contains("decoder") || filename_lower.contains("prompt_encoder"))
+                    && filename_lower.ends_with(".onnx")
+                {
+                    decoder = Some(path);
+                }
+            }
+        }
+
+        match (vision_encoder, decoder) {
+            (Some(vision), Some(dec)) => Ok((vision, dec)),
+            _ => Err(anyhow!("Missing SAM3 model files (need vision_encoder*.onnx and *decoder*.onnx)")),
+        }
+    }
+
+    /// Check if SAM3 model is downloaded
+    pub async fn is_sam3_downloaded(&self) -> Result<bool> {
+        match self.get_sam3_model_paths().await {
+            Ok(_) => Ok(true),
+            Err(_) => Ok(false),
+        }
+    }
+
+    /// Download SAM3 model from HuggingFace
+    pub async fn download_sam3_model<F>(
+        &self,
+        progress_callback: F,
+        cancel_flag: Arc<AtomicBool>,
+    ) -> Result<()>
+    where
+        F: Fn(DownloadProgress) + Send + Clone + 'static,
+    {
+        // SAM3 tracker ONNX models from HuggingFace
+        // Using the quantized Q4 models for efficiency
+        // Note: ONNX models use external data files (.onnx_data) for weights
+        let repo = "onnx-community/sam3-tracker-ONNX";
+        let files = vec![
+            ("onnx/vision_encoder_q4.onnx", "vision_encoder_q4.onnx"),
+            ("onnx/vision_encoder_q4.onnx_data", "vision_encoder_q4.onnx_data"),
+            ("onnx/prompt_encoder_mask_decoder_q4.onnx", "prompt_encoder_mask_decoder_q4.onnx"),
+            ("onnx/prompt_encoder_mask_decoder_q4.onnx_data", "prompt_encoder_mask_decoder_q4.onnx_data"),
+        ];
+
+        let sam3_dir = self.models_dir.join(Self::SAM3_MODEL_ID);
+        fs::create_dir_all(&sam3_dir).await?;
+
+        for (remote_path, local_name) in files {
+            if cancel_flag.load(Ordering::SeqCst) {
+                return Err(anyhow!("Download cancelled"));
+            }
+
+            let url = format!(
+                "https://huggingface.co/{}/resolve/main/{}",
+                repo, remote_path
+            );
+            let local_path = sam3_dir.join(local_name);
+
+            println!("Downloading {} to {:?}", url, local_path);
+
+            self.download_file(&url, &local_path, local_name, progress_callback.clone(), cancel_flag.clone())
+                .await?;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]