Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 import torch
+import torch.nn.functional as F 
 import numpy as np
 from transformers import CLIPImageProcessor
 from modeling_ast import ASTForRestoration
@@ -12,20 +13,21 @@ from tqdm import tqdm
 # --- 1. 配置 ---
 repo_id = "suncongcong/AST_DeRain" 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-PATCH_SIZE = 256  # 模型期望的输入尺寸
-OVERLAP = 64      # 裁切块之间的重叠区域，可以调整
+PATCH_SIZE = 256
+OVERLAP = 64
 
 print(f"正在使用的设备: {device}")
 
 # --- 2. 加载模型和处理器 ---
 print(f"正在从 '{repo_id}' 加载模型和处理器...")
 processor = CLIPImageProcessor.from_pretrained(repo_id)
-# 注意：我们不再修改处理器的尺寸，因为我们会手动裁切
-print(f"图像处理器加载完成。")
+processor.size = {"height": 256, "width": 256}
+processor.crop_size = {"height": 256, "width": 256}
+print(f"图像处理器尺寸已强制设置为: {processor.size}")
 model = ASTForRestoration.from_pretrained(
     repo_id, 
     trust_remote_code=True
-).to(device).eval() # 设置为评估模式
+).to(device).eval()
 print("✅ 模型加载成功，准备就绪！")
 
 
@@ -33,70 +35,63 @@ print("✅ 模型加载成功，准备就绪！")
 def derain_image_Tiled(input_image: Image.Image, progress=gr.Progress(track_tqdm=True)):
     if input_image is None:
         return None
-        
+
     img = input_image.convert("RGB")
     img_tensor = to_tensor(img).unsqueeze(0).to(device)
     b, c, h, w = img_tensor.shape
 
-    # 创建一个空的画布用于存放结果，和一个用于计算平均值的权重图
     output_canvas = torch.zeros_like(img_tensor).to(device)
     weight_map = torch.zeros_like(img_tensor).to(device)
-    
+
     stride = PATCH_SIZE - OVERLAP
-    
-    # 计算需要裁切的块数，用于进度条
-    total_patches = len(range(0, h, stride)) * len(range(0, w, stride))
-    
-    # 使用tqdm来创建进度条
+
+    # 计算需要裁切的块数
+    h_steps = len(range(0, h, stride))
+    w_steps = len(range(0, w, stride))
+    total_patches = h_steps * w_steps
+
     pbar = tqdm(total=total_patches, desc="正在处理图像块...")
 
     for y in range(0, h, stride):
         for x in range(0, w, stride):
-            # 1. 裁切 (Crop)
             y_end = min(y + PATCH_SIZE, h)
             x_end = min(x + PATCH_SIZE, w)
             patch_in = img_tensor[:, :, y:y_end, x:x_end]
-            
-            # 如果边缘块尺寸不够，进行填充 (padding)
+
             ph, pw = patch_in.shape[2:]
             pad_h = PATCH_SIZE - ph
             pad_w = PATCH_SIZE - pw
-            patch_padded = F.pad(patch_in, (0, pad_w, 0, pad_h), 'reflect')
-            
-            # 2. 推理 (Inference)
+            if pad_h > 0 or pad_w > 0:
+                patch_padded = F.pad(patch_in, (0, pad_w, 0, pad_h), 'reflect')
+            else:
+                patch_padded = patch_in
+
             with torch.no_grad():
-                # 注意：这里我们不再使用processor，因为已经手动处理了
-                # 直接将 (1, 3, 256, 256) 的 tensor 送入模型
                 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]
-            
-            # 3. 合并 (Merge)
-            # 将处理后的块加到输出画布上，并更新权重图
+
             output_canvas[:, :, y:y_end, x:x_end] += patch_out_unpadded
             weight_map[:, :, y:y_end, x:x_end] += 1
-            
-            pbar.update(1) # 更新进度条
+
+            pbar.update(1)
 
     pbar.close()
 
-    # 4. 平均 (Average)
-    # 用输出画布除以权重图，得到重叠区域的平均像素值
     restored_tensor = output_canvas / weight_map
-    
+
     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(
         """
-        #AST 图像去雨模型在线演示 (裁切/合并策略)
+        # 🖼️ AST 图像去雨模型在线演示 (裁切/合并策略)
         上传任意尺寸的带雨图片，模型将会分块处理并拼接成完整的高清输出。
         模型仓库地址: [suncongcong/AST_DeRain](https://huggingface.co/suncongcong/AST_DeRain)
         """
@@ -104,10 +99,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     with gr.Row():
         input_img = gr.Image(type="pil", label="输入带雨图片 (Input Rainy Image)")
         output_img = gr.Image(type="pil", label="输出清晰图片 (Output Deraided Image)")
-        
+
     submit_btn = gr.Button("开始去雨 (Start Deraining)", variant="primary")
-    
-    # 将新的处理函数绑定到按钮
+
     submit_btn.click(fn=derain_image_Tiled, inputs=input_img, outputs=output_img)
-    
+
 demo.launch(server_name="0.0.0.0")