diff --git a/output/3T8A4426_nobg.png b/output/3T8A4426_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/3T8A4426_nobg.png differ diff --git a/output/tmp0kh35900_nobg.png b/output/tmp0kh35900_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmp0kh35900_nobg.png differ diff --git a/output/tmp119tqook_nobg.png b/output/tmp119tqook_nobg.png new file mode 100644 index 0000000..9045f12 Binary files /dev/null and b/output/tmp119tqook_nobg.png differ diff --git a/output/tmp3v4i6y57_nobg.png b/output/tmp3v4i6y57_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmp3v4i6y57_nobg.png differ diff --git a/output/tmp4d66uuvp_nobg.png b/output/tmp4d66uuvp_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmp4d66uuvp_nobg.png differ diff --git a/output/tmp4kv2mhwg_nobg.png b/output/tmp4kv2mhwg_nobg.png new file mode 100644 index 0000000..e0875fd Binary files /dev/null and b/output/tmp4kv2mhwg_nobg.png differ diff --git a/output/tmp7e7ckapa_nobg.png b/output/tmp7e7ckapa_nobg.png new file mode 100644 index 0000000..9045f12 Binary files /dev/null and b/output/tmp7e7ckapa_nobg.png differ diff --git a/output/tmp972za3v4_nobg.png b/output/tmp972za3v4_nobg.png new file mode 100644 index 0000000..9045f12 Binary files /dev/null and b/output/tmp972za3v4_nobg.png differ diff --git a/output/tmp_v9_rdjz_nobg.png b/output/tmp_v9_rdjz_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmp_v9_rdjz_nobg.png differ diff --git a/output/tmpb2icryq6_nobg.png b/output/tmpb2icryq6_nobg.png new file mode 100644 index 0000000..9328fd1 Binary files /dev/null and b/output/tmpb2icryq6_nobg.png differ diff --git a/output/tmpb4585ssd_nobg.png b/output/tmpb4585ssd_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmpb4585ssd_nobg.png differ diff --git a/output/tmpdyscjyi5_nobg.png b/output/tmpdyscjyi5_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmpdyscjyi5_nobg.png differ diff --git a/output/tmpfc5ytkx0_nobg.png b/output/tmpfc5ytkx0_nobg.png new file mode 100644 index 0000000..9045f12 Binary files /dev/null and b/output/tmpfc5ytkx0_nobg.png differ diff --git a/output/tmpisrqzyuh_nobg.png b/output/tmpisrqzyuh_nobg.png new file mode 100644 index 0000000..1f47adb Binary files /dev/null and b/output/tmpisrqzyuh_nobg.png differ diff --git a/output/tmpktvlg1nn_nobg.png b/output/tmpktvlg1nn_nobg.png new file mode 100644 index 0000000..3819c55 Binary files /dev/null and b/output/tmpktvlg1nn_nobg.png differ diff --git a/output/tmppdj3mnpp_nobg.png b/output/tmppdj3mnpp_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmppdj3mnpp_nobg.png differ diff --git a/output/tmpsrxe2mao_nobg.png b/output/tmpsrxe2mao_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmpsrxe2mao_nobg.png differ diff --git a/output/tmptv1q37s9_nobg.png b/output/tmptv1q37s9_nobg.png new file mode 100644 index 0000000..91ceee8 Binary files /dev/null and b/output/tmptv1q37s9_nobg.png differ diff --git a/output/tmpxrzisy8i_nobg.png b/output/tmpxrzisy8i_nobg.png new file mode 100644 index 0000000..47408df Binary files /dev/null and b/output/tmpxrzisy8i_nobg.png differ diff --git a/output/tmpycn_udag_nobg.png b/output/tmpycn_udag_nobg.png new file mode 100644 index 0000000..9328fd1 Binary files /dev/null and b/output/tmpycn_udag_nobg.png differ diff --git a/output/tmpym2v1swb_nobg.png b/output/tmpym2v1swb_nobg.png new file mode 100644 index 0000000..bef75db Binary files /dev/null and b/output/tmpym2v1swb_nobg.png differ diff --git a/output/tmpz2kzshzf_nobg.png b/output/tmpz2kzshzf_nobg.png new file mode 100644 index 0000000..2b89660 Binary files /dev/null and b/output/tmpz2kzshzf_nobg.png differ diff --git a/task/rmbg2.py b/task/rmbg2.py index 03f165f..bf209ee 100644 --- a/task/rmbg2.py +++ b/task/rmbg2.py @@ -16,11 +16,8 @@ import gc warnings.filterwarnings("ignore", category=UserWarning) warnings.filterwarnings("ignore", category=FutureWarning) -# 设置torch精度和优化 +# 设置torch精度 torch.set_float32_matmul_precision("high") -torch.backends.cudnn.benchmark = True # 启用cuDNN自动调优 -torch.backends.cuda.matmul.allow_tf32 = True # 允许使用TF32 -torch.backends.cudnn.allow_tf32 = True # 允许cuDNN使用TF32 # 示例图片URL image_url = "http://test001.jingrow.com/files/3T8A4426.JPG" @@ -49,26 +46,31 @@ def process_image(image, model, device): """处理图像，移除背景""" image_size = image.size # 转换图像 + t0 = time.time() transform_image = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) + input_images = transform_image(image).unsqueeze(0).to(device) + print(f"图像预处理耗时: {time.time()-t0:.3f}秒") # 使用半精度加速 + t0 = time.time() with torch.cuda.amp.autocast(): - input_images = transform_image(image).unsqueeze(0).to(device) - # 预测 with torch.no_grad(): preds = model(input_images)[-1].sigmoid().cpu() + print(f"模型推理耗时: {time.time()-t0:.3f}秒") # 处理预测结果 + t0 = time.time() pred = preds[0].squeeze() pred_pil = transforms.ToPILImage()(pred) mask = pred_pil.resize(image_size) # 添加透明通道 image.putalpha(mask) + print(f"后处理耗时: {time.time()-t0:.3f}秒") return image def remove_background(image_path, output_dir=None, model=None, device=None):