# model.py - 完整修复版本，确保与原始测试脚本处理逻辑完全一致 import yaml import torch import math import numpy as np import torch.nn.functional as F from PIL import Image from io import BytesIO import cv2 from basicsr.models import create_model from basicsr.models.archs.HINT_arch import HINT from basicsr.utils.options import parse from skimage import img_as_ubyte import gc import os # 强制使用 CPU，以匹配 Hugging Face Space 的典型环境 os.environ['CUDA_VISIBLE_DEVICES'] = '' torch.backends.cudnn.enabled = False # 根据需要调整线程数，以优化在CPU上的性能 torch.set_num_threads(2) torch.set_num_interop_threads(1) device = torch.device('cpu') # 全局模型变量 dehazing_model = None deraining_model = None desnowing_model = None denoising_blind_model = None enhancement_synthetic_model = None enhancement_real_model = None def splitimage(imgtensor, crop_size=128, overlap_size=64): """原始切块函数 - 与测试脚本完全一致""" _, C, H, W = imgtensor.shape hstarts = [x for x in range(0, H, crop_size - overlap_size)] while hstarts and hstarts[-1] + crop_size >= H: hstarts.pop() hstarts.append(H - crop_size) wstarts = [x for x in range(0, W, crop_size - overlap_size)] while wstarts and wstarts[-1] + crop_size >= W: wstarts.pop() wstarts.append(W - crop_size) starts = [] split_data = [] for hs in hstarts: for ws in wstarts: cimgdata = imgtensor[:, :, hs:hs + crop_size, ws:ws + crop_size] starts.append((hs, ws)) split_data.append(cimgdata) return split_data, starts def get_scoremap(H, W, C, B=1, is_mean=True): """原始权重图生成函数 - 与测试脚本完全一致""" center_h = H / 2 center_w = W / 2 score = torch.ones((B, C, H, W)) if not is_mean: for h in range(H): for w in range(W): score[:, :, h, w] = 1.0 / (math.sqrt((h - center_h) ** 2 + (w - center_w) ** 2 + 1e-6)) return score def mergeimage(split_data, starts, crop_size=128, resolution=(1, 3, 128, 128)): """原始合并函数 - 与测试脚本完全一致""" B, C, H, W = resolution[0], resolution[1], resolution[2], resolution[3] tot_score = torch.zeros((B, C, H, W)) merge_img = torch.zeros((B, C, H, W)) scoremap = get_scoremap(crop_size, crop_size, C, B=B, is_mean=True) for simg, cstart in zip(split_data, starts): hs, ws = cstart merge_img[:, :, hs:hs + crop_size, ws:ws + crop_size] += scoremap * simg tot_score[:, :, hs:hs + crop_size, ws:ws + crop_size] += scoremap merge_img = merge_img / tot_score return merge_img def load_hint_model(config_path, weight_path): """加载HINT模型 - 逻辑与测试脚本完全一致""" try: if not os.path.exists(config_path) or not os.path.exists(weight_path): print(f"HINT模型文件不存在: config='{config_path}', weights='{weight_path}'") return None try: from yaml import CLoader as Loader except ImportError: from yaml import Loader x = yaml.load(open(config_path, mode='r'), Loader=Loader) # 与测试脚本完全一致的网络配置处理 network_config = x['network_g'] model_type = network_config.pop('type') model = HINT(**network_config) # 与测试脚本一致的权重加载方式 checkpoint = torch.load(weight_path, map_location=device) if 'params' in checkpoint: model.load_state_dict(checkpoint['params']) else: model.load_state_dict(checkpoint) model.eval() model = model.to(device) print(f"HINT模型加载成功: {weight_path}") return model except Exception as e: print(f"HINT模型加载失败: {e}") return None def load_basicsr_model(config_path, weight_path): """加载BasicSR模型 - 与测试脚本完全一致""" try: if not os.path.exists(config_path) or not os.path.exists(weight_path): print(f"BasicSR模型文件不存在: config='{config_path}', weights='{weight_path}'") return None opt = parse(config_path, is_train=False) opt['dist'] = False opt['num_gpu'] = 0 opt['device'] = 'cpu' model = create_model(opt).net_g checkpoint = torch.load(weight_path, map_location=device) # 与测试脚本完全一致的权重加载逻辑 try: model.load_state_dict(checkpoint['params']) except: new_checkpoint = {} for k in checkpoint['params']: new_checkpoint['module.' + k] = checkpoint['params'][k] model.load_state_dict(new_checkpoint) model.eval() model = model.to(device) print(f"BasicSR模型加载成功: {weight_path}") return model except Exception as e: print(f"BasicSR模型加载失败: {e}") return None def init(): """初始化所有模型""" global dehazing_model, deraining_model, desnowing_model, denoising_blind_model global enhancement_synthetic_model, enhancement_real_model print("开始加载所有模型...") dehazing_model = load_hint_model("./options/RealDehazing_HINT.yml", "./models/Dehazing.pth") deraining_model = load_hint_model("./options/Deraining_HINT_syn_rain100L.yml", "./models/Rain100L_HINT.pth") desnowing_model = load_hint_model("./options/Desnow_snow100k_HINT.yml", "./models/snow100k.pth") denoising_blind_model = load_hint_model("./options/GaussianColorDenoising_HINT.yml", "./models/net_g_latest_blind.pth") enhancement_synthetic_model = load_basicsr_model("./options/HINT_LOL_v2_synthetic.yml", "./models/lolv2Syn.pth") enhancement_real_model = load_basicsr_model("./options/HINT_LOL_v2_real.yml", "./models/lolv2Real.pth") gc.collect() print("所有模型加载完成。") def preprocess_image(img, task=None, sigma=None): """预处理图像 - 完全按照测试脚本的处理逻辑""" # 转换为numpy数组 if isinstance(img, Image.Image): arr = np.array(img) else: arr = img # 对于去噪任务，完全按照原始代码的处理逻辑 if task == 'denoising_blind' and sigma is not None and sigma > 0: # 步骤1：归一化为 float32 (与原始代码一致) img_normalized = np.float32(arr) / 255.0 # 步骤2：设置随机种子 (与原始代码一致) np.random.seed(seed=0) # 步骤3：添加高斯噪声 (与原始代码完全一致) img_normalized += np.random.normal(0, sigma / 255., img_normalized.shape) # 步骤4：转换为tensor (与原始代码一致: permute(2,0,1)) tensor = torch.from_numpy(img_normalized).permute(2, 0, 1).unsqueeze(0) else: # 其他任务的常规处理 - 完全按照测试脚本逻辑 # 测试脚本: img = np.float32(utils.load_img(inp_path)) / 255. # 测试脚本: img = torch.from_numpy(img).permute(2, 0, 1) arr = np.float32(arr) / 255.0 tensor = torch.from_numpy(arr).permute(2, 0, 1).unsqueeze(0) return tensor def inference(body: bytes, task_type: str, sigma: int = None, GT_mean: bool = False, target_image: bytes = None) -> bytes: """ 统一推理接口 - 完全按照测试脚本的处理逻辑 """ model_map = { "dehazing": dehazing_model, "deraining": deraining_model, "desnowing": desnowing_model, "denoising_blind": denoising_blind_model, "enhancement_synthetic": enhancement_synthetic_model, "enhancement_real": enhancement_real_model } model = model_map.get(task_type) if model is None: raise Exception(f"模型 '{task_type}' 未加载或不存在。") img = Image.open(BytesIO(body)).convert("RGB") input_tensor = preprocess_image(img, task_type, sigma).to(device) orig_h, orig_w = input_tensor.shape[2], input_tensor.shape[3] # 如果需要GT校正，加载目标图像 target_np = None if GT_mean and target_image is not None: target_img = Image.open(BytesIO(target_image)).convert("RGB") target_np = np.array(target_img).astype(np.float32) / 255.0 try: with torch.no_grad(): B, C, H, W = input_tensor.shape # 低光增强任务 - 按照原始测试脚本逻辑处理 if task_type in ["enhancement_synthetic", "enhancement_real"]: factor = 4 # 与测试脚本一致 # 按照原始代码的padding逻辑 H_pad = ((orig_h + factor) // factor) * factor W_pad = ((orig_w + factor) // factor) * factor padh = H_pad - orig_h if orig_h % factor != 0 else 0 padw = W_pad - orig_w if orig_w % factor != 0 else 0 padded_input = F.pad(input_tensor, (0, padw, 0, padh), 'reflect') # 模型推理 restored_tensor = model(padded_input) # 去除padding，恢复原始尺寸 restored_tensor = restored_tensor[:, :, :orig_h, :orig_w] # 去雾任务 - 按照测试脚本逻辑处理 elif task_type == "dehazing": factor = 8 # 按照测试脚本的padding逻辑 H_pad = ((orig_h + factor) // factor) * factor W_pad = ((orig_w + factor) // factor) * factor padh = H_pad - orig_h if orig_h % factor != 0 else 0 padw = W_pad - orig_w if orig_w % factor != 0 else 0 padded_input = F.pad(input_tensor, (0, padw, 0, padh), 'reflect') # 使用测试脚本的切块参数 corp_size_arg = 256 overlap_size_arg = 64 B, C, H_pad, W_pad = padded_input.shape split_data, starts = splitimage(padded_input, crop_size=corp_size_arg, overlap_size=overlap_size_arg) processed_splits = [] for data in split_data: processed_splits.append(model(data).cpu()) restored_tensor = mergeimage(processed_splits, starts, crop_size=corp_size_arg, resolution=(B, C, H_pad, W_pad)) # 去除padding，恢复原始尺寸 restored_tensor = restored_tensor[:, :, :orig_h, :orig_w] # 除雪任务 - 需要切块处理 elif task_type == "desnowing": factor = 8 crop_size = 256 overlap_size = 128 H_pad = ((orig_h + factor) // factor) * factor W_pad = ((orig_w + factor) // factor) * factor padh = H_pad - orig_h if orig_h % factor != 0 else 0 padw = W_pad - orig_w if orig_w % factor != 0 else 0 padded_input = F.pad(input_tensor, (0, padw, 0, padh), 'reflect') split_data, starts = splitimage(padded_input, crop_size=crop_size, overlap_size=overlap_size) processed_splits = [model(s) for s in split_data] restored_tensor = mergeimage(processed_splits, starts, crop_size=crop_size, resolution=(1, 3, H_pad, W_pad)) # 去除padding，恢复原始尺寸 restored_tensor = restored_tensor[:, :, :orig_h, :orig_w] # 去雨和降噪任务 - 全图处理 else: # deraining, denoising_blind factor = 8 H_pad = ((orig_h + factor) // factor) * factor W_pad = ((orig_w + factor) // factor) * factor padh = H_pad - orig_h if orig_h % factor != 0 else 0 padw = W_pad - orig_w if orig_w % factor != 0 else 0 padded_input = F.pad(input_tensor, (0, padw, 0, padh), 'reflect') restored_tensor = model(padded_input) # 去除padding，恢复原始尺寸 restored_tensor = restored_tensor[:, :, :orig_h, :orig_w] # 转换为numpy数组 restored_np = torch.clamp(restored_tensor, 0, 1).cpu().detach().permute(0, 2, 3, 1).squeeze(0).numpy() # GT均值校正 - 完全按照原始代码逻辑 if GT_mean and target_np is not None: # 按照原始代码的逻辑进行GT均值校正 # 原始代码: mean_restored = cv2.cvtColor(restored.astype(np.float32), cv2.COLOR_BGR2GRAY).mean() # 原始代码: mean_target = cv2.cvtColor(target.astype(np.float32), cv2.COLOR_BGR2GRAY).mean() # 原始代码: restored = np.clip(restored * (mean_target / mean_restored), 0, 1) mean_restored = cv2.cvtColor(restored_np.astype(np.float32), cv2.COLOR_RGB2GRAY).mean() mean_target = cv2.cvtColor(target_np.astype(np.float32), cv2.COLOR_RGB2GRAY).mean() restored_np = np.clip(restored_np * (mean_target / mean_restored), 0, 1) # 转换为8位图像格式 result_np = img_as_ubyte(restored_np) out_img = Image.fromarray(result_np) buf = BytesIO() out_img.save(buf, format="PNG") gc.collect() return buf.getvalue() except Exception as e: print(f"在任务 '{task_type}' 中推理失败: {e}") import traceback traceback.print_exc() raise e # --- 导出的推理函数 --- def dehazing_inference(body: bytes) -> bytes: return inference(body, "dehazing") def deraining_inference(body: bytes) -> bytes: return inference(body, "deraining") def desnowing_inference(body: bytes) -> bytes: return inference(body, "desnowing") def denoising_blind_inference(body: bytes, sigma: int = 25) -> bytes: return inference(body, "denoising_blind", sigma=sigma) def enhancement_synthetic_inference(body: bytes, GT_mean: bool = False, target_image: bytes = None) -> bytes: """低光增强（合成）推理""" return inference(body, "enhancement_synthetic", GT_mean=GT_mean, target_image=target_image) def enhancement_real_inference(body: bytes, GT_mean: bool = False, target_image: bytes = None) -> bytes: """低光增强（真实）推理""" return inference(body, "enhancement_real", GT_mean=GT_mean, target_image=target_image) # 为了保持向后兼容性，提供简化的调用方式 def enhancement_synthetic_inference_simple(body: bytes) -> bytes: """简化的低光增强（合成）推理，不使用GT校正""" return inference(body, "enhancement_synthetic") def enhancement_real_inference_simple(body: bytes) -> bytes: """简化的低光增强（真实）推理，不使用GT校正""" return inference(body, "enhancement_real") # --- 辅助函数 --- def get_model_status(): """获取模型加载状态""" return { "dehazing_model_loaded": dehazing_model is not None, "deraining_model_loaded": deraining_model is not None, "desnowing_model_loaded": desnowing_model is not None, "denoising_blind_model_loaded": denoising_blind_model is not None, "enhancement_synthetic_model_loaded": enhancement_synthetic_model is not None, "enhancement_real_model_loaded": enhancement_real_model is not None } def get_available_tasks(): """获取可用的任务列表""" return ['dehazing', 'deraining', 'desnowing', 'denoising_blind', 'enhancement_synthetic', 'enhancement_real'] def get_task_description(task): """获取任务描述""" descriptions = { 'dehazing': '去除图像中的雾霾效果', 'deraining': '去除图像中的雨线', 'desnowing': '去除图像中的雪花', 'denoising_blind': '通用降噪', 'enhancement_synthetic': '低光增强 (合成数据训练)', 'enhancement_real': '低光增强 (真实数据训练)' } return descriptions.get(task, '未知任务') # --- 测试函数 --- def test_model_loading(): """测试模型加载情况""" print("测试模型加载情况:") status = get_model_status() for model_name, loaded in status.items(): print(f" {model_name}: {'✓' if loaded else '✗'}") print() def test_inference(task_type, test_image_path=None): """测试推理功能""" if test_image_path is None: # 创建一个测试图像 test_img = Image.new('RGB', (256, 256), color='red') buf = BytesIO() test_img.save(buf, format='PNG') test_bytes = buf.getvalue() else: with open(test_image_path, 'rb') as f: test_bytes = f.read() try: if task_type == 'dehazing': result = dehazing_inference(test_bytes) elif task_type == 'deraining': result = deraining_inference(test_bytes) elif task_type == 'desnowing': result = desnowing_inference(test_bytes) elif task_type == 'denoising_blind': result = denoising_blind_inference(test_bytes, sigma=25) elif task_type == 'enhancement_synthetic': result = enhancement_synthetic_inference(test_bytes) elif task_type == 'enhancement_real': result = enhancement_real_inference(test_bytes) else: print(f"未知任务类型: {task_type}") return False print(f"✓ {task_type} 推理测试成功，输出大小: {len(result)} bytes") return True except Exception as e: print(f"✗ {task_type} 推理测试失败: {e}") return False if __name__ == "__main__": # 初始化模型 init() # 测试模型加载 test_model_loading() # 测试所有任务的推理 print("测试推理功能:") for task in get_available_tasks(): test_inference(task) print("\n所有测试完成！")