优化rmbg2.py

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):