Update app.py

app.py

@@ -20,7 +20,7 @@ MODEL_IDS = {
 EXAMPLE_IMAGES = {
     "去雨痕 (Derain)": [["derain_example1.png"], ["derain_example2.png"], ["derain_example3.png"]],
     "去雨滴 (Deraindrop)": [["deraindrop_example1.png"], ["deraindrop_example2.png"], ["deraindrop_example3.png"]],
-    "去雾 (Dehaze)": [["dehaze_example1.jpg"],["dehaze_example2.jpg"],["dehaze_example3.jpg"]] 
+    "去雾 (Dehaze)": [["dehaze_example1.jpg"],["dehaze_example2.jpg"],["dehaze_example3.jpg"]]
 }
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"正在使用的设备: {device}")
@@ -45,41 +45,68 @@ except Exception as e:
     MODELS = {task: load_error_func for task in MODEL_IDS.keys()}
 print("所有模型加载完毕，准备就绪！")
 
-
-# --- 3. 定义统一的处理函数 ---
-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]
-    if isinstance(model, str): raise gr.Error(model)
-
-    print(f"执行任务: {task_name}, 使用 Pad-to-Square 策略")
-    img_tensor = to_tensor(input_image.convert("RGB")).unsqueeze(0)
-    
-    # Pad-to-Square 逻辑
+# --- 3. 定义不同任务的处理函数 ---
+def process_with_pad_to_square(model, img_tensor):
     def expand2square(timg, factor=128.0):
-        _, _, h, w = timg.size(); X = int(math.ceil(max(h, w) / 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)
         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)
-        
     restored_tensor = torch.masked_select(
         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
+    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
+    img_padded = F.pad(img_tensor, (0, w_pad, 0, h_pad), 'replicate')
+    _, _, h_padded, w_padded = img_padded.shape
+    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]:
+            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)
+        
     restored_tensor = torch.clamp(restored_tensor, 0, 1)
-    return to_pil_image(restored_tensor.cpu().squeeze(0))
-
+    restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
+    return restored_image
 
-# --- 4. 创建并启动 Gradio 界面 ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
@@ -94,19 +121,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                     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")
                 
-                # 使用 gr.State 来传递 task_name，这是更稳健的方式
-                task_state = gr.State(task_name)
-                
-                submit_btn.click(fn=process_image, inputs=[input_img, task_state], outputs=output_img)
+                # “提交”按钮的点击事件，保持不变
+                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],
+                        inputs=input_img,
                         outputs=output_img,
-                        fn=lambda img: process_image(img, task_name=task_name),
+                        fn=create_example_fn(task_name), # 关键：为每个循环的 task_name 创建一个独立的函数
                         cache_examples=True
                     )