This commit is contained in:
jhqxxx
2025-12-26 00:33:07 +08:00
+6 -6
View File
@@ -92,12 +92,13 @@ impl RMBG2_0Model {
return Ok(vec![]);
}
// 并行预处理:提取原始尺寸和转换为 tensor
// 并行预处理:提取原始尺寸、RGB 数据和转换为 tensor
let preprocessed: Vec<_> = imgs
.par_iter()
.map(|img| {
let height = img.height();
let width = img.width();
let rgb_img = img.to_rgb8();
let tensor = img_transform_with_resize(
img,
self.h,
@@ -107,17 +108,17 @@ impl RMBG2_0Model {
&self.device,
self.dtype,
);
(img.clone(), height, width, tensor)
(rgb_img, height, width, tensor)
})
.collect();
// 检查预处理是否有错误
let mut tensors = Vec::with_capacity(preprocessed.len());
let mut meta: Vec<_> = Vec::with_capacity(preprocessed.len());
for (img, height, width, tensor_result) in preprocessed {
for (rgb_img, height, width, tensor_result) in preprocessed {
let tensor = tensor_result?;
tensors.push(tensor);
meta.push((img, height, width));
meta.push((rgb_img, height, width));
}
// 批量推理:将所有图片合并为一个 batch
@@ -134,7 +135,7 @@ impl RMBG2_0Model {
let results: Vec<Result<RgbaImage>> = meta
.into_par_iter()
.enumerate()
.map(|(i, (img, height, width))| {
.map(|(i, (rgb_img, height, width))| {
// let rmbg_tensor = batch_output.i(i)?;
let rmbg_tensor = &batch_output[i];
let alpha_img = float_tensor_to_dynamic_image(rmbg_tensor)?;
@@ -142,7 +143,6 @@ impl RMBG2_0Model {
alpha_img.resize_exact(width, height, image::imageops::FilterType::CatmullRom);
let alpha_gray = alpha_img.to_luma8();
let rgb_img = img.to_rgb8();
let rgb_raw = rgb_img.as_raw();
let alpha_raw = alpha_gray.as_raw();
let pixel_count = (width * height) as usize;