SR-Label-SamTool/sam-tool-tauri/src-tauri/src/lib.rs

mod annotate;
mod annotation;
mod coco;
mod decoder;
mod extract;
mod sam;
mod settings;
mod state;
mod user;
mod video;

use decoder::DecoderSession;
use serde::Serialize;
use state::AppState;
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use tauri::{ipc::Response, AppHandle, Emitter, State};
use video::VideoInfo;

#[tauri::command]
fn ping() -> String {
    "pong".into()
}

#[tauri::command]
fn open_video(state: State<'_, AppState>, path: String) -> Result<VideoInfo, String> {
    let p = PathBuf::from(&path);
    if !p.exists() {
        return Err(format!("file not found: {}", p.display()));
    }
    let info = video::probe(&p)?;
    let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
    *guard = Some(DecoderSession::open(p, info.clone()));
    Ok(info)
}

#[tauri::command]
fn close_video(state: State<'_, AppState>) -> Result<(), String> {
    let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
    *guard = None;
    Ok(())
}

#[tauri::command]
fn current_video(state: State<'_, AppState>) -> Result<Option<VideoInfo>, String> {
    let guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
    Ok(guard.as_ref().map(|d| d.info().clone()))
}

#[tauri::command]
fn decode_frame(state: State<'_, AppState>, frame_idx: u64) -> Result<Response, String> {
    let bytes = {
        let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
        let decoder = guard.as_mut().ok_or("no video loaded")?;
        decoder.decode_at(frame_idx)?
    };
    Ok(Response::new(bytes))
}

#[tauri::command]
fn import_coco(path: String) -> Result<coco::ImportedCoco, String> {
    let p = Path::new(&path);
    if !p.exists() {
        return Err(format!("file not found: {}", p.display()));
    }
    coco::import(p)
}

#[tauri::command]
fn set_cache_capacity(state: State<'_, AppState>, capacity: usize) -> Result<usize, String> {
    let clamped = capacity.clamp(decoder::MIN_CACHE_CAPACITY, decoder::MAX_CACHE_CAPACITY);
    let guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
    if let Some(d) = guard.as_ref() {
        d.set_cache_capacity(clamped);
    }
    Ok(clamped)
}

#[tauri::command]
fn extract_frame(
    state: State<'_, AppState>,
    frame_idx: u64,
    out_dir: String,
    shapes: Vec<coco::Shape>,
    author: String,
) -> Result<(), String> {
    let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
    let decoder = guard.as_mut().ok_or("no video loaded")?;
    extract::extract_one(decoder, frame_idx, Path::new(&out_dir), &shapes, &author)
}

#[tauri::command]
fn undo_extract(frame_idx: u64, out_dir: String) -> Result<(), String> {
    extract::undo_one(frame_idx, Path::new(&out_dir))
}

#[derive(Serialize, Clone)]
struct BulkProgress {
    done: u64,
    total: u64,
}

#[tauri::command]
fn bulk_extract(
    state: State<'_, AppState>,
    app: AppHandle,
    start: u64,
    end: u64,
    out_dir: String,
    shapes_by_frame: HashMap<String, Vec<coco::Shape>>,
    author: String,
) -> Result<u64, String> {
    if end < start {
        return Err("end must be >= start".into());
    }
    let mut guard = state.decoder.lock().map_err(|e| format!("lock: {e}"))?;
    let decoder = guard.as_mut().ok_or("no video loaded")?;
    let dir = PathBuf::from(&out_dir);
    let total = end - start + 1;
    let empty: Vec<coco::Shape> = Vec::new();
    let mut done: u64 = 0;
    for idx in start..=end {
        let shapes = shapes_by_frame
            .get(&idx.to_string())
            .unwrap_or(&empty);
        extract::extract_one(decoder, idx, &dir, shapes, &author)?;
        done += 1;
        if done == 1 || done % 5 == 0 || done == total {
            let _ = app.emit("bulk-extract-progress", BulkProgress { done, total });
        }
    }
    Ok(done)
}

#[tauri::command]
fn bulk_undo(
    app: AppHandle,
    start: u64,
    end: u64,
    out_dir: String,
) -> Result<u64, String> {
    if end < start {
        return Err("end must be >= start".into());
    }
    let dir = PathBuf::from(&out_dir);
    let total = end - start + 1;
    let mut done: u64 = 0;
    for idx in start..=end {
        extract::undo_one(idx, &dir)?;
        done += 1;
        if done == 1 || done % 20 == 0 || done == total {
            let _ = app.emit("bulk-undo-progress", BulkProgress { done, total });
        }
    }
    Ok(done)
}

#[tauri::command]
fn list_extracted_frames(out_dir: String) -> Result<Vec<u64>, String> {
    extract::list_extracted(Path::new(&out_dir))
}

#[tauri::command]
fn write_classes_txt(out_dir: String, categories: Vec<coco::Category>) -> Result<(), String> {
    extract::write_classes_txt(Path::new(&out_dir), &categories)
}

#[tauri::command]
fn list_folder_images(folder: String) -> Result<Vec<annotate::ImageEntry>, String> {
    annotate::list_images(Path::new(&folder))
}

#[tauri::command]
fn read_image_bytes(path: String) -> Result<Response, String> {
    let bytes = annotate::read_image_bytes(Path::new(&path))?;
    Ok(Response::new(bytes))
}

#[tauri::command]
fn load_image_annotations(
    image_path: String,
) -> Result<Option<annotation::ImageAnnotations>, String> {
    annotate::load_image_annotations(Path::new(&image_path))
}

#[tauri::command]
fn load_classes(folder: String) -> Result<Vec<String>, String> {
    annotate::load_classes(Path::new(&folder))
}

#[tauri::command]
fn append_class(folder: String, name: String) -> Result<Vec<String>, String> {
    annotate::append_class(Path::new(&folder), &name)
}

#[tauri::command]
fn delete_images(paths: Vec<String>) -> Result<u32, String> {
    annotate::delete_images(&paths)
}

#[tauri::command]
async fn sam_health_check(url: String) -> Result<bool, String> {
    sam::health_check(&url).await
}

#[tauri::command]
async fn sam_list_models(url: String) -> Result<sam::ModelList, String> {
    sam::list_models(&url).await
}

#[tauri::command]
async fn sam_segment_bbox(
    url: String,
    image_path: String,
    bbox: [f64; 4],
    image_width: u32,
    image_height: u32,
    model: Option<String>,
) -> Result<Vec<[f64; 2]>, String> {
    sam::segment_bbox(
        &url,
        Path::new(&image_path),
        bbox,
        image_width,
        image_height,
        model,
    )
    .await
}

#[tauri::command]
fn save_image_annotations(
    image_path: String,
    payload: annotation::ImageAnnotations,
) -> Result<(), String> {
    let p = Path::new(&image_path);
    let json_path = p.with_extension("json");
    if payload.annotations.is_empty() {
        if json_path.exists() {
            std::fs::remove_file(&json_path)
                .map_err(|e| format!("remove {}: {e}", json_path.display()))?;
        }
        return Ok(());
    }
    // Backfill source_video from folder-level source.json when the payload
    // doesn't carry it — handles annotations authored via AnnotateMode that
    // never went through extract_one.
    let mut payload = payload;
    if payload.source_video.is_none() {
        if let Some(parent) = p.parent() {
            if let Some(meta) = extract::read_folder_source(parent) {
                payload.source_video = Some(meta.source_video);
            }
        }
    }
    let text = serde_json::to_string_pretty(&payload)
        .map_err(|e| format!("serialize: {e}"))?;
    std::fs::write(&json_path, text)
        .map_err(|e| format!("write {}: {e}", json_path.display()))
}

#[tauri::command]
fn read_folder_source(folder: String) -> Option<String> {
    extract::read_folder_source(Path::new(&folder)).map(|s| s.source_video)
}

#[cfg_attr(mobile, tauri::mobile_entry_point)]
pub fn run() {
    tauri::Builder::default()
        .plugin(tauri_plugin_dialog::init())
        .plugin(tauri_plugin_fs::init())
        .manage(AppState::new())
        .invoke_handler(tauri::generate_handler![
            ping,
            user::current_user,
            user::list_known_users,
            user::set_current_user,
            user::clear_current_user,
            open_video,
            close_video,
            current_video,
            decode_frame,
            set_cache_capacity,
            import_coco,
            extract_frame,
            undo_extract,
            bulk_extract,
            bulk_undo,
            list_extracted_frames,
            write_classes_txt,
            list_folder_images,
            read_image_bytes,
            load_image_annotations,
            load_classes,
            save_image_annotations,
            append_class,
            delete_images,
            read_folder_source,
            settings::get_settings,
            settings::save_settings,
            sam_health_check,
            sam_list_models,
            sam_segment_bbox,
        ])
        .run(tauri::generate_context!())
        .expect("error while running tauri application");
}