diff --git a/src/models/rmbg2_0/generate.rs b/src/models/rmbg2_0/generate.rs index 961b90c..db5db3d 100644 --- a/src/models/rmbg2_0/generate.rs +++ b/src/models/rmbg2_0/generate.rs @@ -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> = 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;