refactor: 为RMBG2.0测试方法添加cfg(test)属性,优化extract_image_url函数避免clone
This commit is contained in:
@@ -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();
|
||||||
|
|||||||
@@ -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);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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);
|
||||||
|
|||||||
Reference in New Issue
Block a user