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" } 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"]] } 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}' 加载成功。") 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("所有模型加载完毕，准备就绪！") # --- 3. 定义不同任务的处理函数 --- def process_with_pad_to_square(model, img_tensor): """将图片填充为正方形后进行处理，适用于去雨/去雨滴任务。""" def expand2square(timg, factor=128.0): # factor: 模型的网络结构要求输入的尺寸最好是该值的整数倍 _, _, h, w = timg.size() X = int(math.ceil(max(h, w) / factor) * factor) # 确保创建的张量在正确的设备上 img_padded = torch.zeros(1, 3, X, X, device=timg.device, dtype=timg.dtype) mask = torch.zeros(1, 1, X, X, device=timg.device, dtype=timg.dtype) pad_h = (X - h) // 2 pad_w = (X - w) // 2 img_padded[:, :, pad_h:pad_h + h, pad_w:pad_w + w] = timg mask[:, :, pad_h:pad_h + h, pad_w:pad_w + 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) # 确保 mask 和 restored_padded 在同一设备上 mask_bool = mask.bool().to(restored_padded.device) 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 = 1152 # 每个图块的尺寸 OVERLAP = 384 # 图块之间的重叠区域大小，以避免边缘效应 b, c, h_orig, w_orig = img_tensor.shape stride = CROP_SIZE - OVERLAP # 计算需要填充的尺寸 h_pad = (stride - (h_orig - OVERLAP) % stride) % stride if h_orig > OVERLAP else 0 w_pad = (stride - (w_orig - OVERLAP) % stride) % stride if w_orig > OVERLAP else 0 img_padded = F.pad(img_tensor, (0, w_pad, 0, h_pad), 'replicate') b, c, h_padded, w_padded = img_padded.shape # 使用CPU来存储最终结果，避免占用大量显存 output_canvas = torch.zeros((b, c, h_padded, w_padded), device='cpu') weight_map = torch.zeros_like(output_canvas) h_steps_range = range(0, h_padded - OVERLAP, stride) if h_padded > OVERLAP else [0] w_steps_range = range(0, w_padded - OVERLAP, stride) if w_padded > OVERLAP else [0] # 使用Gradio的进度条 for y in progress.tqdm(h_steps_range, desc="正在分块去雾..."): for x in w_steps_range: # 确保切片范围正确 y_end = min(y + CROP_SIZE, h_padded) x_end = min(x + CROP_SIZE, w_padded) patch_in = img_padded[:, :, y:y_end, x:x_end] with torch.no_grad(): patch_out = model(patch_in.to(device)).cpu() output_canvas[:, :, y:y_end, x:x_end] += patch_out weight_map[:, :, y:y_end, x:x_end] += 1 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: gr.Warning("请输入一张图片！") return None # 增加完整的运行时错误捕获 try: model = MODELS[task_name] print(f"已选择任务: {task_name}, 使用模型: {MODEL_IDS[task_name]}") # 检查模型是否加载成功 if not isinstance(model, torch.nn.Module): model() # 如果加载失败，这里会触发 load_error_func 并抛出异常 img = input_image.convert("RGB") img_tensor = to_tensor(img).unsqueeze(0) 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 except Exception as e: print(f"处理图片时发生错误: {e}") # 在UI上给用户一个清晰的错误提示 gr.Error(f"处理失败！错误信息: {e}") # 返回原始图像，而不是空着或保留上一次的结果 return input_image # --- 4. Gradio UI --- def create_task_tab(task_name: str): """动态创建每个任务的UI选项卡。""" 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") # ✨ 修正后的处理函数 ✨ # 这个函数只接收它能从 inputs 中得到的 `img` 参数。 def specific_process_fn(img): # 调用 process_image 时不传递 progress 参数， # 从而让 process_image 自动使用其函数定义中的默认值: progress=gr.Progress(...) return process_image(img, task_name) # click 事件的 inputs 列表只有一个元素，对应 specific_process_fn 的 img 参数 submit_btn.click(fn=specific_process_fn, inputs=[input_img], outputs=output_img) if EXAMPLE_IMAGES.get(task_name): gr.Examples( examples=EXAMPLE_IMAGES.get(task_name, []), inputs=input_img, outputs=output_img, fn=specific_process_fn, # 复用上面为按钮创建的处理函数 cache_examples=True, ) # 创建应用主界面 with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 🖼️ 多功能图像复原工具 (AST 模型) 请选择一个任务，然后上传对应的图片或点击下方的示例图片进行处理。 """ ) with gr.Tabs(): for task_name in MODEL_IDS.keys(): create_task_tab(task_name) # 调用函数为每个任务创建Tab # 启动应用 demo.launch(server_name="0.0.0.0")