Update app.py

app.py

@@ -9,114 +9,126 @@ import requests
 from io import BytesIO
 from torchvision.transforms.functional import to_pil_image, to_tensor
 from tqdm import tqdm
+import math
 
 # --- 1. 配置 ---
-# 使用您提供的准确的模型仓库ID
 MODEL_IDS = {
     "去雨痕 (Derain)": "Suncongcong/AST_DeRain",
     "去雨滴 (Deraindrop)": "Suncongcong/AST_DeRainDrop",
     "去雾 (Dehaze)": "Suncongcong/AST_Dehazing"
 }
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
-PATCH_SIZE = 256
-OVERLAP = 64
-
 print(f"正在使用的设备: {device}")
 
 # --- 2. 加载所有模型和处理器 ---
 MODELS = {}
 PROCESSOR = None
-
 print("正在加载所有模型和处理器...")
-# 使用 try-except 来增加鲁棒性
 try:
     for task_name, repo_id in MODEL_IDS.items():
         print(f"正在加载模型: {task_name} ({repo_id})")
         if PROCESSOR is None:
             PROCESSOR = CLIPImageProcessor.from_pretrained(repo_id)
             print("✅ 处理器加载成功。")
-        
-        model = ASTForRestoration.from_pretrained(
-            repo_id, 
-            trust_remote_code=True
-        ).to(device).eval()
+        model = ASTForRestoration.from_pretrained(repo_id, trust_remote_code=True).to(device).eval()
         MODELS[task_name] = model
         print(f"✅ 模型 '{task_name}' 加载成功。")
 except Exception as e:
     print(f"加载模型时出错: {e}")
-    # 创建一个占位符函数，以便在模型加载失败时 Gradio 仍能启动并显示错误
     def load_error_func(*args, **kwargs):
         raise gr.Error(f"模型加载失败! 错误: {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
+# --- 3. 定义不同任务的处理函数 ---
+
+# 策略一：用于去雨痕、去雨滴的 "Pad-to-Square" 函数
+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)
+        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)
     
-    # 根据传入的任务名称，选择对应的模型
-    model = MODELS[task_name]
-    print(f"已选择任务: {task_name}, 使用模型: {MODEL_IDS[task_name]}")
+    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)
     
-    # 检查模型是否加载成功
-    if not isinstance(model, torch.nn.Module):
-        model() # 这会触发上面定义的错误函数
-
-    img = input_image.convert("RGB")
-    img_tensor = to_tensor(img).unsqueeze(0).to(device)
-    b, c, h, w = img_tensor.shape
+    return restored_tensor
 
-    output_canvas = torch.zeros_like(img_tensor).to(device)
-    weight_map = torch.zeros_like(img_tensor).to(device)
+# 策略二：用于去雾的 "Crop-and-Merge" 函数
+def process_with_dehaze_tiling(model, img_tensor, progress):
+    PATCH_SIZE = 1152
+    OVERLAP = 384
     
+    b, c, h, w = img_tensor.shape
+    output_canvas = torch.zeros((b, c, h, w), device='cpu') # 合并时在CPU上操作
+    weight_map = torch.zeros_like(output_canvas)
     stride = PATCH_SIZE - OVERLAP
     
-    h_steps = len(range(0, h, stride))
-    w_steps = len(range(0, w, stride))
+    # 填充以确保图像尺寸至少为一个Patch大小
+    pad_h = max(0, PATCH_SIZE - h)
+    pad_w = max(0, PATCH_SIZE - w)
+    if pad_h > 0 or pad_w > 0:
+        img_tensor = F.pad(img_tensor, (0, pad_w, 0, pad_h), 'reflect')
+        _, _, h, w = img_tensor.shape
+
+    h_steps = len(range(0, h - OVERLAP, stride)) if h > OVERLAP else 1
+    w_steps = len(range(0, w - OVERLAP, stride)) if w > OVERLAP else 1
     total_patches = h_steps * w_steps
-    
-    pbar = tqdm(total=total_patches, desc=f"正在执行 {task_name}...")
+    pbar = tqdm(total=total_patches, desc="正在处理去雾...")
 
-    for y in range(0, h, stride):
-        for x in range(0, w, stride):
-            y_end = min(y + PATCH_SIZE, h)
-            x_end = min(x + PATCH_SIZE, w)
-            patch_in = img_tensor[:, :, y:y_end, x:x_end]
-            
-            ph, pw = patch_in.shape[2:]
-            pad_h = PATCH_SIZE - ph
-            pad_w = PATCH_SIZE - pw
-            if pad_h > 0 or pad_w > 0:
-                patch_padded = F.pad(patch_in, (0, pad_w, 0, pad_h), 'replicate')
-            else:
-                patch_padded = patch_in
-            
+    for y in range(0, h - OVERLAP, stride) if h > OVERLAP else [0]:
+        for x in range(0, w - OVERLAP, stride) if w > OVERLAP else [0]:
+            patch_in = img_tensor[:, :, y:y+PATCH_SIZE, x:x+PATCH_SIZE]
             with torch.no_grad():
-                outputs = model(patch_padded)
-            
-            patch_out = outputs[0] if isinstance(outputs, tuple) else outputs
-            patch_out = torch.clamp(patch_out, 0, 1)
-
-            patch_out_unpadded = patch_out[:, :, :ph, :pw]
-            
-            output_canvas[:, :, y:y_end, x:x_end] += patch_out_unpadded
-            weight_map[:, :, y:y_end, x:x_end] += 1
+                patch_out = model(patch_in.to(device)).cpu()
             
+            output_canvas[:, :, y:y+PATCH_SIZE, x:x+PATCH_SIZE] += patch_out
+            weight_map[:, :, y:y+PATCH_SIZE, x:x+PATCH_SIZE] += 1
             pbar.update(1)
-
     pbar.close()
 
-    restored_tensor = output_canvas / weight_map
-    restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
+    restored_tensor = (output_canvas / torch.clamp(weight_map, min=1))
+    # 裁切掉为了计算而额外填充的部分
+    final_h, final_w = resolution[2:]
+    restored_tensor = restored_tensor[:, :, :final_h, :final_w]
+    return restored_tensor
+
+# 主调度函数
+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)
+    
+    global resolution
+    resolution = img_tensor.shape
+
+    if task_name == "去雾 (Dehaze)":
+        restored_tensor = process_with_dehaze_tiling(model, img_tensor, progress)
+    else: # 去雨痕和去雨滴使用 Pad-to-Square 策略
+        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
 
-# --- 4. 创建并启动带选项卡的 Gradio 界面 ---
+
+# --- 4. 创建并启动 Gradio 界面 (无变化) ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(
         """
@@ -124,23 +136,14 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         请选择一个任务，然后上传对应的图片进行处理。
         """
     )
-    
     with gr.Tabs():
-        # 根据 MODEL_IDS 字典自动创建选项卡
         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
-                )
+                submit_btn.click(fn=process_image, inputs=[input_img, task_id_box], outputs=output_img)
 
 demo.launch(server_name="0.0.0.0")