diff --git a/src/models/rmbg2_0/generate.rs b/src/models/rmbg2_0/generate.rs index d355d7c..6306db4 100644 --- a/src/models/rmbg2_0/generate.rs +++ b/src/models/rmbg2_0/generate.rs @@ -45,30 +45,36 @@ impl RMBG2_0 { }) } + #[cfg(test)] pub fn h(&self) -> u32 { self.h } + #[cfg(test)] pub fn w(&self) -> u32 { self.w } + #[cfg(test)] pub fn img_mean(&self) -> &Tensor { &self.img_mean } + #[cfg(test)] pub fn img_std(&self) -> &Tensor { &self.img_std } + #[cfg(test)] pub fn device(&self) -> &Device { &self.device } - + #[cfg(test)] pub fn dtype(&self) -> DType { self.dtype } + #[cfg(test)] pub fn model(&self) -> &BiRefNet { &self.model } @@ -118,8 +124,8 @@ impl RMBG2_0 { .map(|(i, (img, height, width))| { let rmbg_tensor = batch_output.i(i)?; let alpha_img = float_tensor_to_dynamic_image(&rmbg_tensor)?; - let alpha_img = alpha_img - .resize_exact(width, height, image::imageops::FilterType::CatmullRom); + let alpha_img = + alpha_img.resize_exact(width, height, image::imageops::FilterType::CatmullRom); let alpha_gray = alpha_img.to_luma8(); let rgb_img = img.to_rgb8(); diff --git a/src/utils/img_utils.rs b/src/utils/img_utils.rs index 0aee160..9a3e2a6 100644 --- a/src/utils/img_utils.rs +++ b/src/utils/img_utils.rs @@ -79,16 +79,15 @@ pub fn get_image(file: &str) -> Result { Err(anyhow!("get image from message failed".to_string())) } -pub fn extract_image_url(mes: &ChatCompletionParameters) -> Vec { +pub fn extract_image_url(mes: &ChatCompletionParameters) -> Vec<&String> { let mut img_vec = Vec::new(); - // 使用引用避免 clone for chat_mes in &mes.messages { if let ChatMessage::User { content, .. } = chat_mes && let ChatMessageContent::ContentPart(part_vec) = content { for part in part_vec { if let ChatMessageContentPart::Image(img_part) = part { - img_vec.push(img_part.image_url.url.clone()); + img_vec.push(&img_part.image_url.url); } } } diff --git a/tests/test_rmbg2_0_perf.rs b/tests/test_rmbg2_0_perf.rs index 8793f3c..57a9019 100644 --- a/tests/test_rmbg2_0_perf.rs +++ b/tests/test_rmbg2_0_perf.rs @@ -21,7 +21,7 @@ fn test_pixel_combine_performance() -> Result<()> { image::Luma([((x + y) % 256) as u8]) }); - let iterations = 50; + let iterations = 10; println!("=== 像素组合性能测试 ==="); println!("图片尺寸: {}x{}", width, height); @@ -195,7 +195,7 @@ fn test_postprocess_parallel_vs_serial() -> Result<()> { }) .collect(); - let iterations = 20; + let iterations = 10; println!("=== 后处理阶段性能测试(纯图像操作)==="); println!("图片数量: {}", num_images);