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("正在加载所有模型和处理器...")
@@ -45,67 +45,141 @@ except Exception as e:
     MODELS = {task: load_error_func for task in MODEL_IDS.keys()}
 print("所有模型加载完毕，准备就绪！")
 
+
 # --- 3. 定义不同任务的处理函数 ---
+
 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).type_as(timg)
-        mask = torch.zeros(1, 1, X, X).type_as(timg)
-        img_padded[:, :, ((X - h) // 2):((X - h) // 2 + h), ((X - w) // 2):((X - w) // 2 + w)] = timg
-        mask[:, :, ((X - h) // 2):((X - h) // 2 + h), ((X - w) // 2):((X - w) // 2 + w)].fill_(1)
+        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)
+        
+        pad_h = (X - h) // 2
+        pad_w = (X - w) // 2
+        img_padded[:, :, pad_h:pad_h + h, pad_w:pad_w + w] = timg
+        mask[:, :, pad_h:pad_h + h, pad_w:pad_w + w].fill_(1)
         return img_padded, mask
+        
     original_h, original_w = img_tensor.shape[2:]
     padded_input, mask = expand2square(img_tensor.to(device), factor=128.0)
+    
     with torch.no_grad():
         restored_padded = model(padded_input)
+        
+    # ✨ 优化点: 确保 mask 和 restored_padded 在同一设备上
+    mask_bool = mask.bool().to(restored_padded.device)
+    
     restored_tensor = torch.masked_select(
-        restored_padded, mask.bool()
+        restored_padded, mask_bool
     ).reshape(1, 3, original_h, original_w)
+    
     return restored_tensor
 
+
 def process_with_dehaze_tiling(model, img_tensor, progress):
-    CROP_SIZE, OVERLAP = 1152, 384
+    """使用重叠分块策略处理图像，适用于去雾任务。"""
+    # ✨ 优化点: 将“魔法数字”定义为常量并添加注释
+    CROP_SIZE = 1152  # 每个图块的尺寸
+    OVERLAP = 384   # 图块之间的重叠区域大小，以避免边缘效应
+    
     b, c, h_orig, w_orig = img_tensor.shape
     stride = CROP_SIZE - OVERLAP
-    h_pad = (stride - (h_orig - OVERLAP) % stride) % stride
-    w_pad = (stride - (w_orig - OVERLAP) % stride) % stride
+
+    # 计算需要填充的尺寸
+    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
+
     img_padded = F.pad(img_tensor, (0, w_pad, 0, h_pad), 'replicate')
-    _, _, h_padded, w_padded = img_padded.shape
+    b, c, h_padded, w_padded = img_padded.shape
+
+    # 使用CPU来存储最终结果，避免占用大量显存
     output_canvas = torch.zeros((b, c, h_padded, w_padded), device='cpu')
     weight_map = torch.zeros_like(output_canvas)
-    h_steps = len(range(0, h_padded - OVERLAP, stride)) if h_padded > OVERLAP else 1
-    w_steps = len(range(0, w_padded - OVERLAP, stride)) if w_padded > OVERLAP else 1
-    pbar = tqdm(total=h_steps * w_steps, desc=f"正在执行去雾...")
-    for y in range(0, h_padded - OVERLAP, stride) if h_padded > OVERLAP else [0]:
-        for x in range(0, w_padded - OVERLAP, stride) if w_padded > OVERLAP else [0]:
+    
+    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]
+    
+    total_steps = len(h_steps_range) * len(w_steps_range)
+    
+    # ✨ 优化点: 使用Gradio的进度条，而不是手动的tqdm
+    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]
             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
-            pbar.update(1)
-    pbar.close()
+            
     restored_padded_tensor = output_canvas / torch.clamp(weight_map, min=1)
     return restored_padded_tensor[:, :, :h_orig, :w_orig]
 
+
 def process_image(input_image: Image.Image, task_name: str, progress=gr.Progress(track_tqdm=True)):
-    if input_image is None: return None
-    model = MODELS[task_name]
-    print(f"已选择任务: {task_name}, 使用模型: {MODEL_IDS[task_name]}")
-    if not isinstance(model, torch.nn.Module): model()
-    img = input_image.convert("RGB")
-    img_tensor = to_tensor(img).unsqueeze(0)
-    
-    # 关键修正：在 process_image 函数内部也进行判断
-    if task_name == "去雾 (Dehaze)":
-        restored_tensor = process_with_dehaze_tiling(model, img_tensor, progress)
-    else: 
-        restored_tensor = process_with_pad_to_square(model, img_tensor)
+    """主处理函数，根据任务名分派到不同的处理流程。"""
+    if input_image is None:
+        gr.Warning("请输入一张图片！")
+        return None
+        
+    # ✨ 优化点: 增加完整的运行时错误捕获
+    try:
+        model = MODELS[task_name]
+        print(f"已选择任务: {task_name}, 使用模型: {MODEL_IDS[task_name]}")
+        
+        # 检查模型是否加载成功
+        if not isinstance(model, torch.nn.Module):
+            model() # 如果加载失败，这里会触发 load_error_func 并抛出异常
+        
+        img = input_image.convert("RGB")
+        img_tensor = to_tensor(img).unsqueeze(0)
+        
+        if task_name == "去雾 (Dehaze)":
+            restored_tensor = process_with_dehaze_tiling(model, img_tensor, progress)
+        else:
+            restored_tensor = process_with_pad_to_square(model, img_tensor)
+            
+        restored_tensor = torch.clamp(restored_tensor, 0, 1)
+        restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
+        return restored_image
+        
+    except Exception as e:
+        print(f"处理图片时发生错误: {e}")
+        # 在UI上给用户一个清晰的错误提示
+        gr.Error(f"处理失败！错误信息: {e}")
+        # 返回原始图像，而不是空着或保留上一次的结果
+        return input_image
+
+
+# --- 4. Gradio UI ---
+
+# ✨ 优化点 3: 将创建Tab的逻辑封装成函数，使UI代码更干净
+def create_task_tab(task_name: str):
+    with gr.TabItem(task_name, id=task_name):
+        with gr.Row():
+            input_img = gr.Image(type="pil", label=f"输入图片 (Input for {task_name})")
+            output_img = gr.Image(type="pil", label="输出图片 (Output)")
+        
+        submit_btn = gr.Button("开始处理 (Process)", variant="primary")
+        
+        # ✨ 优化点 1: 创建一个处理函数，它已经“知道”自己的任务名
+        # 这样就不再需要隐藏的 Textbox 来传递 task_name
+        def specific_process_fn(img, prog):
+            return process_image(img, task_name, prog)
+
+        submit_btn.click(fn=specific_process_fn, inputs=[input_img], outputs=output_img)
         
-    restored_tensor = torch.clamp(restored_tensor, 0, 1)
-    restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
-    return restored_image
+        if EXAMPLE_IMAGES.get(task_name):
+            gr.Examples(
+                examples=EXAMPLE_IMAGES.get(task_name, []),
+                inputs=input_img,
+                outputs=output_img,
+                fn=specific_process_fn, # 复用上面为按钮创建的处理函数
+                cache_examples=True,
+            )
 
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
@@ -116,30 +190,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
     with gr.Tabs():
         for task_name in MODEL_IDS.keys():
-            with gr.TabItem(task_name, id=task_name):
-                with gr.Row():
-                    input_img = gr.Image(type="pil", label=f"输入图片 (Input for {task_name})")
-                    output_img = gr.Image(type="pil", label="输出图片 (Output)")
-                
-                task_id_box = gr.Textbox(task_name, visible=False)
-                submit_btn = gr.Button("开始处理 (Process)", variant="primary")
-                
-                # “提交”按钮的点击事件，保持不变
-                submit_btn.click(fn=process_image, inputs=[input_img, task_id_box], outputs=output_img)
-                
-                # --- 最终修正 ---
-                # 重新构造 lambda 函数，确保它总是传递正确的 task_name
-                # 我们不再依赖外部的 task_name 变量，而是直接使用在循环中定义的那个
-                def create_example_fn(current_task_name):
-                    return lambda img, prog: process_image(img, current_task_name, prog)
-
-                if EXAMPLE_IMAGES.get(task_name):
-                    gr.Examples(
-                        examples=EXAMPLE_IMAGES.get(task_name, []),
-                        inputs=input_img,
-                        outputs=output_img,
-                        fn=create_example_fn(task_name), # 关键：为每个循环的 task_name 创建一个独立的函数
-                        cache_examples=True
-                    )
+            create_task_tab(task_name) # 调用函数创建每个Tab
 
 demo.launch(server_name="0.0.0.0")