Genarrative/server-rs/crates/platform-matting/examples/segment_smoke.rs

//! 通用抠图冒烟验证：本地图片 → AuthorizeFileUpload 临时对象 → SegmentCommonImage → 下载结果。
//!
//! 运行（在 server-rs 目录下）：
//! cargo run -p platform-matting --example segment_smoke -- "C:\path\to\input.png"
//!
//! 依赖仓库根目录 .env / .env.local / .env.secrets.local 中的 AK/SK
//!（`GENARRATIVE_ALIYUN_MATTING_*` 或 `ALIBABA_CLOUD_ACCESS_KEY_*`）；
//! 可选 `GENARRATIVE_ALIYUN_MATTING_ENDPOINT`。临时 OSS bucket/endpoint 由
//! AuthorizeFileUpload 动态下发，不依赖自有 OSS 配置。

use std::path::{Path, PathBuf};

use platform_matting::{
    DEFAULT_IMAGESEG_ENDPOINT, MattingClient, MattingConfig, SegmentCommonImageRequest,
    SegmentReturnForm,
};

fn load_env_files() {
    // example 的工作目录通常是 server-rs，.env 都在仓库根目录。
    for candidate in [
        ".env",
        ".env.local",
        ".env.secrets.local",
        "../.env",
        "../.env.local",
        "../.env.secrets.local",
    ] {
        if Path::new(candidate).exists() {
            let _ = dotenvy::from_path_override(candidate);
        }
    }
}

#[tokio::main]
async fn main() {
    load_env_files();

    let input_path = std::env::args().nth(1).unwrap_or_else(|| {
        eprintln!("用法：cargo run -p platform-matting --example segment_smoke -- <图片路径>");
        std::process::exit(1);
    });
    let input_bytes = std::fs::read(&input_path)
        .unwrap_or_else(|error| panic!("读取测试图片失败（{input_path}）：{error}"));
    println!(
        "[1/5] 已读取测试图片：{input_path}（{} 字节）",
        input_bytes.len()
    );

    // SegmentCommonImage 要求分辨率低于 2000x2000，超限先等比缩小。
    const MAX_EDGE: u32 = 1999;
    let decoded = image::load_from_memory(&input_bytes).expect("测试图片应可解码");
    let input_bytes = if decoded.width() > MAX_EDGE || decoded.height() > MAX_EDGE {
        let resized = decoded.resize(MAX_EDGE, MAX_EDGE, image::imageops::FilterType::CatmullRom);
        let mut buffer = std::io::Cursor::new(Vec::new());
        resized
            .write_to(&mut buffer, image::ImageFormat::Png)
            .expect("缩放后的图片应可编码为 PNG");
        println!(
            "      分辨率 {}x{} 超限，已缩放到 {}x{}",
            decoded.width(),
            decoded.height(),
            resized.width(),
            resized.height()
        );
        buffer.into_inner()
    } else {
        input_bytes
    };

    let http_client = reqwest::Client::new();

    // --- 调用通用抠图 ---
    // key 优先级与 api-server 配置保持一致：抠图专用 → 官方 SDK 标准命名。
    let (matting_key_id, matting_key_secret) = [
        (
            "GENARRATIVE_ALIYUN_MATTING_ACCESS_KEY_ID",
            "GENARRATIVE_ALIYUN_MATTING_ACCESS_KEY_SECRET",
        ),
        (
            "ALIBABA_CLOUD_ACCESS_KEY_ID",
            "ALIBABA_CLOUD_ACCESS_KEY_SECRET",
        ),
    ]
    .iter()
    .find_map(|(id_name, secret_name)| {
        let id = std::env::var(id_name).ok()?;
        let secret = std::env::var(secret_name).ok()?;
        if id.trim().is_empty() || secret.trim().is_empty() {
            return None;
        }
        println!("      使用 {id_name} 调用抠图服务");
        Some((id, secret))
    })
    .expect("未找到可用的抠图 AccessKey 环境变量");
    let matting_config = MattingConfig::new(
        std::env::var("GENARRATIVE_ALIYUN_MATTING_ENDPOINT")
            .unwrap_or_else(|_| DEFAULT_IMAGESEG_ENDPOINT.to_string()),
        matting_key_id,
        matting_key_secret,
    )
    .expect("抠图配置应有效");
    let matting_client = MattingClient::new(matting_config).expect("抠图客户端应可构建");

    // 非上海地域输入按新版官方 SDK 的 AdvanceRequest 口径申请单对象 Policy 后上传。
    let temp_url = matting_client
        .upload_temp_image(input_bytes, "segment-input.png", "image/png")
        .await
        .expect("上传 AuthorizeFileUpload 临时对象应成功");
    println!("[2/5] 已上传 AuthorizeFileUpload 临时对象");
    println!("[3/5] 输入图 URL host：{}", host_of(&temp_url));

    let result = matting_client
        .segment_common_image(SegmentCommonImageRequest {
            image_url: temp_url,
            return_form: Some(SegmentReturnForm::Crop),
        })
        .await;
    let result = match result {
        Ok(result) => result,
        Err(error) => {
            eprintln!("[4/5] 通用抠图调用失败：{error}");
            std::process::exit(1);
        }
    };
    println!(
        "[4/5] 通用抠图成功，RequestId={}",
        result.request_id.as_deref().unwrap_or("unknown")
    );

    // --- 下载结果 ---
    let output_bytes = http_client
        .get(&result.image_url)
        .send()
        .await
        .expect("下载抠图结果应成功")
        .error_for_status()
        .expect("抠图结果 URL 应返回 200")
        .bytes()
        .await
        .expect("读取抠图结果字节应成功");

    let output_path = build_output_path(&input_path);
    std::fs::write(&output_path, &output_bytes).expect("写出抠图结果应成功");
    println!(
        "[5/5] 抠图结果已保存：{}（{} 字节）",
        output_path.display(),
        output_bytes.len()
    );
}

fn build_output_path(input_path: &str) -> PathBuf {
    let input = Path::new(input_path);
    let stem = input
        .file_stem()
        .and_then(|value| value.to_str())
        .unwrap_or("segment-output");
    input.with_file_name(format!("{stem}-matting.png"))
}

fn host_of(url: &str) -> String {
    url.split("//")
        .nth(1)
        .and_then(|rest| rest.split('/').next())
        .unwrap_or("unknown")
        .to_string()
}