Update app.py

app.py

@@ -1,283 +1,142 @@
 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)": [["dehaze_example1.jpg"],["dehaze_example2.jpg"],["dehaze_example3.jpg"]]
-
+    "去雨痕 (Derain)": [["derain_example1.png"], ["derain_example2.png"], ["derain_example3.png"]],
+    "去雨滴 (Deraindrop)": [["deraindrop_example1.png"], ["deraindrop_example2.png"], ["deraindrop_example3.png"]],
+    "去雾 (Dehaze)": []
 }
-
 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)
-
-    pbar.close()
-
-    restored_padded_tensor = output_canvas / torch.clamp(weight_map, min=1)
-
-    return restored_padded_tensor[:, :, :h_orig, :w_orig]
-
-
+    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
-
-
+    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)")
-
-                
-
-                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
-
-                    )
-
-
+    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")