Update app.py

app.py

@@ -1,51 +1,113 @@
 import gradio as gr
 import torch
+import numpy as np
 from transformers import CLIPImageProcessor
 from modeling_ast import ASTForRestoration
 from PIL import Image
 import requests
 from io import BytesIO
-from torchvision.transforms.functional import to_pil_image
+from torchvision.transforms.functional import to_pil_image, to_tensor
+from tqdm import tqdm
 
-# --- 1. 配置模型和设备 ---
+# --- 1. 配置 ---
 repo_id = "suncongcong/AST_DeRain" 
 device = "cuda" if torch.cuda.is_available() else "cpu"
+PATCH_SIZE = 256  # 模型期望的输入尺寸
+OVERLAP = 64      # 裁切块之间的重叠区域，可以调整
+
 print(f"正在使用的设备: {device}")
 
-# --- 2. 加载模型和图像处理器 ---
+# --- 2. 加载模型和处理器 ---
 print(f"正在从 '{repo_id}' 加载模型和处理器...")
 processor = CLIPImageProcessor.from_pretrained(repo_id)
-processor.size = {"height": 256, "width": 256}
-processor.crop_size = {"height": 256, "width": 256}
-print(f"图像处理器尺寸已强制设置为: {processor.size}")
+# 注意：我们不再修改处理器的尺寸，因为我们会手动裁切
+print(f"图像处理器加载完成。")
 model = ASTForRestoration.from_pretrained(
     repo_id, 
     trust_remote_code=True
-).to(device)
+).to(device).eval() # 设置为评估模式
 print("✅ 模型加载成功，准备就绪！")
 
-# --- 3. 定义核心处理函数 ---
-def derain_image(input_image: Image.Image):
-    if input_image is None: return None
-    image = input_image.convert("RGB")
-    inputs = processor(images=image, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    restored_tensor = outputs[0] if isinstance(outputs, tuple) else outputs
-    restored_tensor = torch.clamp(restored_tensor, 0, 1)
+
+# --- 3. 定义“裁切-推理-合并”的核心处理函数 ---
+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来创建进度条
+    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)
+            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.close()
+
+    # 4. 平均 (Average)
+    # 用输出画布除以权重图，得到重叠区域的平均像素值
+    restored_tensor = output_canvas / weight_map
+    
     restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
+    
     return restored_image
 
 # --- 4. 创建并启动 Gradio 界面 ---
-print("正在创建 Gradio Interface...")
-demo = gr.Interface(
-    fn=derain_image,
-    inputs=gr.Image(type="pil", label="输入带雨图片 (Input Rainy Image)"),
-    outputs=gr.Image(type="pil", label="输出清晰图片 (Output Deraided Image)"),
-    title="AST 图像去雨模型在线演示",
-    description="上传一张带雨的图片，模型将会自动去除雨水痕迹。模型仓库地址: [suncongcong/AST_DeRain](https://huggingface.co/suncongcong/AST_DeRain)"
-)
-
-# --- 最终修正：添加 server_name 参数以适应容器环境 ---
-print("正在启动 Demo...")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        #AST 图像去雨模型在线演示 (裁切/合并策略)
+        上传任意尺寸的带雨图片，模型将会分块处理并拼接成完整的高清输出。
+        模型仓库地址: [suncongcong/AST_DeRain](https://huggingface.co/suncongcong/AST_DeRain)
+        """
+    )
+    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")