Update app.py

app.py

@@ -1,127 +1,283 @@
 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
+
 from PIL import Image
+
 import requests
+
 from io import BytesIO
+
 from torchvision.transforms.functional import to_pil_image, to_tensor
+
 from tqdm import tqdm
+
 import math
 
+
+
 # --- 1. 配置 ---
+
 MODEL_IDS = {
-    "去雨痕 (Derain)": "Suncongcong/AST_DeRain",
-    "去雨滴 (Deraindrop)": "Suncongcong/AST_DeRainDrop",
-    "去雾 (Dehaze)": "Suncongcong/AST_Dehazing"
+
+    "去雨痕 (Derain)": "Suncongcong/AST_DeRain",
+
+    "去雨滴 (Deraindrop)": "Suncongcong/AST_DeRainDrop",
+
+    "去雾 (Dehaze)": "Suncongcong/AST_Dehazing"
+
 }
+
 EXAMPLE_IMAGES = {
-    "去雨痕 (Derain)": [["derain_example1.png"], ["derain_example2.png"], ["derain_example3.png"]],
-    "去雨滴 (Deraindrop)": [["deraindrop_example1.png"], ["deraindrop_example2.png"], ["deraindrop_example3.png"]],
-    "去雾 (Dehaze)": []
+
+    "去雨痕 (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"]]
+
 }
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
+
 print(f"正在使用的设备: {device}")
 
+
+
 # --- 2. 加载所有模型和处理器 ---
+
 MODELS = {}
+
 PROCESSOR = None
+
 print("正在加载所有模型和处理器...")
+
 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()
-        MODELS[task_name] = model
-        print(f"✅ 模型 '{task_name}' 加载成功。")
+
+    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()
+
+        MODELS[task_name] = model
+
+        print(f"✅ 模型 '{task_name}' 加载成功。")
+
 except Exception as e:
-    print(f"加载模型时出错: {e}")
-    def load_error_func(*args, **kwargs):
-        raise gr.Error(f"模型加载失败! 错误: {e}")
-    MODELS = {task: load_error_func for task in MODEL_IDS.keys()}
+
+    print(f"加载模型时出错: {e}")
+
+    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_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)
-    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 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)
+
+    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)
-            # --- 最终修正：只有在 progress 对象存在时才更新它 ---
-            if progress:
-                progress(pbar.n / pbar.total, desc=pbar.desc)
-    pbar.close()
-    restored_padded_tensor = output_canvas / torch.clamp(weight_map, min=1)
-    return restored_padded_tensor[:, :, :h_orig, :w_orig]
-
-# 主调度函数
-def run_task(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)
-    img_tensor = to_tensor(input_image.convert("RGB")).unsqueeze(0)
-    print(f"执行任务: {task_name}")
-    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))
-
-# --- 4. 创建并启动 Gradio 界面 ---
+
+    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)
+
+    restored_image = to_pil_image(restored_tensor.cpu().squeeze(0))
+
+    return restored_image
+
+
+
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # 🖼️ 多功能图像复原工具 (AST 模型)
-        请选择一个任务，然后上传对应的图片或点击下方的示例图片进行处理。
-        """
-    )
-    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)")
-                submit_btn = gr.Button("开始处理 (Process)", variant="primary")
-                submit_btn.click(fn=run_task, inputs=[input_img, gr.State(task_name)], outputs=output_img)
-                if EXAMPLE_IMAGES.get(task_name):
-                    gr.Examples(
-                        examples=EXAMPLE_IMAGES.get(task_name, []),
-                        inputs=input_img,
-                        outputs=output_img,
-                        fn=lambda img: run_task(img, task_name=task_name), # 示例点击时不传递 progress
-                        cache_examples=True
-                    )
+
+    gr.Markdown(
+
+        """
+
+        # 🖼️ 多功能图像复原工具 (AST 模型)
+
+        请选择一个任务，然后上传对应的图片或点击下方的示例图片进行处理。
+
+        """
+
+    )
+
+    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)
+
+                
+
+                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
+
+                    )
+
+
 
 demo.launch(server_name="0.0.0.0")