update
This commit is contained in:
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user