Update app.py

app.py

@@ -11,7 +11,7 @@ from torchvision.transforms.functional import to_pil_image, to_tensor
 from tqdm import tqdm
 import math
 
-# --- 1. 配置 (无变化) ---
+# --- 1. 配置 ---
 MODEL_IDS = {
     "去雨痕 (Derain)": "Suncongcong/AST_DeRain",
     "去雨滴 (Deraindrop)": "Suncongcong/AST_DeRainDrop",
@@ -25,7 +25,7 @@ EXAMPLE_IMAGES = {
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"正在使用的设备: {device}")
 
-# --- 2. 加载所有模型和处理器 (无变化) ---
+# --- 2. 加载所有模型和处理器 ---
 MODELS = {}
 PROCESSOR = None
 print("正在加载所有模型和处理器...")
@@ -51,10 +51,10 @@ print("所有模型加载完毕，准备就绪！")
 def process_with_pad_to_square(model, img_tensor):
     """将图片填充为正方形后进行处理，适用于去雨/去雨滴任务。"""
     def expand2square(timg, factor=128.0):
-        # ✨ 优化点: 增加注释，解释 factor 的作用
         # factor: 模型的网络结构要求输入的尺寸最好是该值的整数倍
         _, _, h, w = timg.size()
         X = int(math.ceil(max(h, w) / factor) * factor)
+        # 确保创建的张量在正确的设备上
         img_padded = torch.zeros(1, 3, X, X, device=timg.device, dtype=timg.dtype)
         mask = torch.zeros(1, 1, X, X, device=timg.device, dtype=timg.dtype)
         
@@ -70,7 +70,7 @@ def process_with_pad_to_square(model, img_tensor):
     with torch.no_grad():
         restored_padded = model(padded_input)
         
-    # ✨ 优化点: 确保 mask 和 restored_padded 在同一设备上
+    # 确保 mask 和 restored_padded 在同一设备上
     mask_bool = mask.bool().to(restored_padded.device)
     
     restored_tensor = torch.masked_select(
@@ -82,7 +82,7 @@ def process_with_pad_to_square(model, img_tensor):
 
 def process_with_dehaze_tiling(model, img_tensor, progress):
     """使用重叠分块策略处理图像，适用于去雾任务。"""
-    # ✨ 优化点: 将“魔法数字”定义为常量并添加注释
+    # 将“魔法数字”定义为常量并添加注释
     CROP_SIZE = 1152  # 每个图块的尺寸
     OVERLAP = 384   # 图块之间的重叠区域大小，以避免边缘效应
     
@@ -90,8 +90,8 @@ def process_with_dehaze_tiling(model, img_tensor, progress):
     stride = CROP_SIZE - OVERLAP
 
     # 计算需要填充的尺寸
-    h_pad = (stride - (h_orig - OVERLAP) % stride) % stride if h_orig > CROP_SIZE else 0
-    w_pad = (stride - (w_orig - OVERLAP) % stride) % stride if w_orig > CROP_SIZE else 0
+    h_pad = (stride - (h_orig - OVERLAP) % stride) % stride if h_orig > OVERLAP else 0
+    w_pad = (stride - (w_orig - OVERLAP) % stride) % stride if w_orig > OVERLAP else 0
 
     img_padded = F.pad(img_tensor, (0, w_pad, 0, h_pad), 'replicate')
     b, c, h_padded, w_padded = img_padded.shape
@@ -100,20 +100,22 @@ def process_with_dehaze_tiling(model, img_tensor, progress):
     output_canvas = torch.zeros((b, c, h_padded, w_padded), device='cpu')
     weight_map = torch.zeros_like(output_canvas)
     
-    h_steps_range = range(0, h_padded - CROP_SIZE + 1, stride) if h_padded > CROP_SIZE else [0]
-    w_steps_range = range(0, w_padded - CROP_SIZE + 1, stride) if w_padded > CROP_SIZE else [0]
+    h_steps_range = range(0, h_padded - OVERLAP, stride) if h_padded > OVERLAP else [0]
+    w_steps_range = range(0, w_padded - OVERLAP, stride) if w_padded > OVERLAP else [0]
     
-    total_steps = len(h_steps_range) * len(w_steps_range)
-    
-    # ✨ 优化点: 使用Gradio的进度条，而不是手动的tqdm
+    # 使用Gradio的进度条
     for y in progress.tqdm(h_steps_range, desc="正在分块去雾..."):
         for x in w_steps_range:
-            patch_in = img_padded[:, :, y:y+CROP_SIZE, x:x+CROP_SIZE]
+            # 确保切片范围正确
+            y_end = min(y + CROP_SIZE, h_padded)
+            x_end = min(x + CROP_SIZE, w_padded)
+            patch_in = img_padded[:, :, y:y_end, x:x_end]
+
             with torch.no_grad():
                 patch_out = model(patch_in.to(device)).cpu()
             
-            output_canvas[:, :, y:y+CROP_SIZE, x:x+CROP_SIZE] += patch_out
-            weight_map[:, :, y:y+CROP_SIZE, x:x+CROP_SIZE] += 1
+            output_canvas[:, :, y:y_end, x:x_end] += patch_out
+            weight_map[:, :, y:y_end, x:x_end] += 1
             
     restored_padded_tensor = output_canvas / torch.clamp(weight_map, min=1)
     return restored_padded_tensor[:, :, :h_orig, :w_orig]
@@ -125,7 +127,7 @@ def process_image(input_image: Image.Image, task_name: str, progress=gr.Progress
         gr.Warning("请输入一张图片！")
         return None
         
-    # ✨ 优化点: 增加完整的运行时错误捕获
+    # 增加完整的运行时错误捕获
     try:
         model = MODELS[task_name]
         print(f"已选择任务: {task_name}, 使用模型: {MODEL_IDS[task_name]}")
@@ -156,8 +158,8 @@ def process_image(input_image: Image.Image, task_name: str, progress=gr.Progress
 
 # --- 4. Gradio UI ---
 
-# ✨ 优化点 3: 将创建Tab的逻辑封装成函数，使UI代码更干净
 def create_task_tab(task_name: str):
+    """动态创建每个任务的UI选项卡。"""
     with gr.TabItem(task_name, id=task_name):
         with gr.Row():
             input_img = gr.Image(type="pil", label=f"输入图片 (Input for {task_name})")
@@ -165,11 +167,14 @@ def create_task_tab(task_name: str):
         
         submit_btn = gr.Button("开始处理 (Process)", variant="primary")
         
-        # ✨ 优化点 1: 创建一个处理函数，它已经“知道”自己的任务名
-        # 这样就不再需要隐藏的 Textbox 来传递 task_name
-        def specific_process_fn(img, prog):
-            return process_image(img, task_name, prog)
-
+        # ✨ 修正后的处理函数 ✨
+        # 这个函数只接收它能从 inputs 中得到的 `img` 参数。
+        def specific_process_fn(img):
+            # 调用 process_image 时不传递 progress 参数，
+            # 从而让 process_image 自动使用其函数定义中的默认值: progress=gr.Progress(...)
+            return process_image(img, task_name)
+
+        # click 事件的 inputs 列表只有一个元素，对应 specific_process_fn 的 img 参数
         submit_btn.click(fn=specific_process_fn, inputs=[input_img], outputs=output_img)
         
         if EXAMPLE_IMAGES.get(task_name):
@@ -181,6 +186,7 @@ def create_task_tab(task_name: str):
                 cache_examples=True,
             )
 
+# 创建应用主界面
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
@@ -190,6 +196,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
     with gr.Tabs():
         for task_name in MODEL_IDS.keys():
-            create_task_tab(task_name) # 调用函数创建每个Tab
+            create_task_tab(task_name) # 调用函数为每个任务创建Tab
 
+# 启动应用
 demo.launch(server_name="0.0.0.0")