refactor: 为RMBG2.0测试方法添加cfg(test)属性,优化extract_image_url函数避免clone

This commit is contained in:
michaelbguo
2025-12-25 10:44:32 +08:00
parent f35631157a
commit 7c9261567b
3 changed files with 13 additions and 8 deletions
+9 -3
View File
@@ -45,30 +45,36 @@ impl RMBG2_0 {
}) })
} }
#[cfg(test)]
pub fn h(&self) -> u32 { pub fn h(&self) -> u32 {
self.h self.h
} }
#[cfg(test)]
pub fn w(&self) -> u32 { pub fn w(&self) -> u32 {
self.w self.w
} }
#[cfg(test)]
pub fn img_mean(&self) -> &Tensor { pub fn img_mean(&self) -> &Tensor {
&self.img_mean &self.img_mean
} }
#[cfg(test)]
pub fn img_std(&self) -> &Tensor { pub fn img_std(&self) -> &Tensor {
&self.img_std &self.img_std
} }
#[cfg(test)]
pub fn device(&self) -> &Device { pub fn device(&self) -> &Device {
&self.device &self.device
} }
#[cfg(test)]
pub fn dtype(&self) -> DType { pub fn dtype(&self) -> DType {
self.dtype self.dtype
} }
#[cfg(test)]
pub fn model(&self) -> &BiRefNet { pub fn model(&self) -> &BiRefNet {
&self.model &self.model
} }
@@ -118,8 +124,8 @@ impl RMBG2_0 {
.map(|(i, (img, height, width))| { .map(|(i, (img, height, width))| {
let rmbg_tensor = batch_output.i(i)?; let rmbg_tensor = batch_output.i(i)?;
let alpha_img = float_tensor_to_dynamic_image(&rmbg_tensor)?; let alpha_img = float_tensor_to_dynamic_image(&rmbg_tensor)?;
let alpha_img = alpha_img let alpha_img =
.resize_exact(width, height, image::imageops::FilterType::CatmullRom); alpha_img.resize_exact(width, height, image::imageops::FilterType::CatmullRom);
let alpha_gray = alpha_img.to_luma8(); let alpha_gray = alpha_img.to_luma8();
let rgb_img = img.to_rgb8(); let rgb_img = img.to_rgb8();
+2 -3
View File
@@ -79,16 +79,15 @@ pub fn get_image(file: &str) -> Result<DynamicImage> {
Err(anyhow!("get image from message failed".to_string())) Err(anyhow!("get image from message failed".to_string()))
} }
pub fn extract_image_url(mes: &ChatCompletionParameters) -> Vec<String> { pub fn extract_image_url(mes: &ChatCompletionParameters) -> Vec<&String> {
let mut img_vec = Vec::new(); let mut img_vec = Vec::new();
// 使用引用避免 clone
for chat_mes in &mes.messages { for chat_mes in &mes.messages {
if let ChatMessage::User { content, .. } = chat_mes if let ChatMessage::User { content, .. } = chat_mes
&& let ChatMessageContent::ContentPart(part_vec) = content && let ChatMessageContent::ContentPart(part_vec) = content
{ {
for part in part_vec { for part in part_vec {
if let ChatMessageContentPart::Image(img_part) = part { if let ChatMessageContentPart::Image(img_part) = part {
img_vec.push(img_part.image_url.url.clone()); img_vec.push(&img_part.image_url.url);
} }
} }
} }
+2 -2
View File
@@ -21,7 +21,7 @@ fn test_pixel_combine_performance() -> Result<()> {
image::Luma([((x + y) % 256) as u8]) image::Luma([((x + y) % 256) as u8])
}); });
let iterations = 50; let iterations = 10;
println!("=== 像素组合性能测试 ==="); println!("=== 像素组合性能测试 ===");
println!("图片尺寸: {}x{}", width, height); println!("图片尺寸: {}x{}", width, height);
@@ -195,7 +195,7 @@ fn test_postprocess_parallel_vs_serial() -> Result<()> {
}) })
.collect(); .collect();
let iterations = 20; let iterations = 10;
println!("=== 后处理阶段性能测试(纯图像操作)==="); println!("=== 后处理阶段性能测试(纯图像操作)===");
println!("图片数量: {}", num_images); println!("图片数量: {}", num_images);