Pix2Pix 4通道推理 Demo

# app.py # -*- coding: utf-8 -*- import os, sys, io, uuid, time, traceback from typing import Optional from PIL import Image from flask import Flask, request, jsonify, render_template_string, send_from_directory from werkzeug.utils import secure_filename import torch import torch.nn as nn import torchvision.transforms as T from typing import Optional # ========= 你训练仓库的路径：改成你的实际路径 ========= PIX2PIX_ROOT = "/Users/liutao/Downloads/python/pix2pix_local" sys.path.insert(0, PIX2PIX_ROOT) # 从你的训练仓库导入“网络构造函数” # 官方实现一般是 models/networks.py 里提供 define_G from models.networks import define_G # ---------------- 基本配置 ---------------- WEIGHTS_PATH = "./weights/150_net_G.pth" # 你的权重 UPLOAD_DIR = "./runtime/uploads" OUTPUT_DIR = "./runtime/outputs" os.makedirs(UPLOAD_DIR, exist_ok=True) os.makedirs(OUTPUT_DIR, exist_ok=True) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" IMG_H, IMG_W = 1536, 512 NORM_TO_MINUS1_1 = True # Pix2Pix 通常使用 [-1,1] 归一化 # ---------------- Flask ---------------- app = Flask(__name__) # 放在 app = Flask(__name__) 下面 from flask_cors import CORS CORS(app, resources={r"/api/*": {"origins": "*"}}, supports_credentials=True) # 允许较大上传（比如 50MB） app.config["MAX_CONTENT_LENGTH"] = 50 * 1024 * 1024 # 统一给响应加上允许头（有时浏览器仍会要） @app.after_request def add_cors_headers(resp): resp.headers["Access-Control-Allow-Origin"] = "*" resp.headers["Access-Control-Allow-Methods"] = "GET,POST,OPTIONS" resp.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization" return resp INDEX_HTML = r""" Pix2Pix 4通道推理 Demo

Pix2Pix 4通道（RGB+黑mask, 512×1536）在线推理

当前权重：

上传一张图片（将被缩放到 512×1536，并在后端叠加全黑 mask 作为第4通道）：

输出

""" # ---------------- 模型加载（关键） ---------------- NETG: Optional[nn.Module] = None def build_netG() -> nn.Module: """ 复用你训练时的参数： --netG unet_512 --norm instance --no_dropout 且你是 4通道输入（RGB+mask），3通道输出官方实现 define_G 的典型签名： define_G(input_nc, output_nc, ngf=64, netG='unet_256', norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, gpu_ids=[]) """ netG = define_G( input_nc=4, # 你的输入 RGB+mask = 4 output_nc=3, # 输出 RGB = 3 ngf=64, netG='unet_512', # 训练用的 unet_512 norm='instance', # 训练用的 instance use_dropout=False # 训练时 --no_dropout ) return netG def load_weights_into(net: nn.Module, weight_path: str): state = torch.load(weight_path, map_location="cpu") # 有些保存方式会套一层 for key in ["state_dict", "netG", "model", "module"]: if isinstance(state, dict) and key in state and isinstance(state[key], (dict, torch.nn.modules.module.Module)): # 常见：{'state_dict': {...}} 或 {'netG': state_dict} state = state[key] break # 如果键名带 'module.' 前缀，去掉 if isinstance(state, dict): new_state = {} for k, v in state.items(): nk = k.replace("module.", "") # DataParallel 保存会多这个前缀 new_state[nk] = v state = new_state missing, unexpected = net.load_state_dict(state, strict=False) print(f"[load_state_dict] missing: {missing}, unexpected: {unexpected}") def init_model(): global NETG netG = build_netG() load_weights_into(netG, WEIGHTS_PATH) netG.eval().to(DEVICE) NETG = netG print(f"[OK] netG loaded on {DEVICE} from {WEIGHTS_PATH}") # ---------------- 预处理/后处理 ---------------- def pil_to_tensor_4ch(img: Image.Image, mask_img: Optional[Image.Image] = None) -> torch.Tensor: # 1) 尺寸对齐：只有不一致才调整 img = img.convert("RGB") if img.size != (IMG_W, IMG_H): # 你的数据就是 512x1536；不一致时才处理：等比按高缩放，再左裁/右补到 512 new_w = round(img.width * IMG_H / img.height) img = img.resize((new_w, IMG_H), Image.BICUBIC) if new_w >= IMG_W: img = img.crop((0, 0, IMG_W, IMG_H)) else: pad = Image.new("RGB", (IMG_W, IMG_H), (0,0,0)) pad.paste(img, (0,0)) img = pad if mask_img is not None: mask_img = mask_img.convert("L") if mask_img.size != (IMG_W, IMG_H): # 对 mask 只能用最近邻，避免灰边 new_w = round(mask_img.width * IMG_H / mask_img.height) mask_img = mask_img.resize((new_w, IMG_H), Image.NEAREST) if new_w >= IMG_W: mask_img = mask_img.crop((0, 0, IMG_W, IMG_H)) else: padm = Image.new("L", (IMG_W, IMG_H), 0) padm.paste(mask_img, (0,0)) mask_img = padm m = T.ToTensor()(mask_img) # [1,H,W], 值∈[0,1] else: m = torch.zeros(1, IMG_H, IMG_W) # 2) 与 test.py 同分布：ToTensor->[0,1] 再 *2-1 x3 = T.ToTensor()(img) # [3,H,W], [0,1] x3 = x3 * 2.0 - 1.0 # [-1,1] 等价于 Normalize(0.5,0.5) x4 = torch.cat([x3, m], dim=0) # [4,H,W] return x4 def tensor_to_pil(y: torch.Tensor) -> Image.Image: if y.dim() == 4: y = y[0] y = y.detach().cpu() if NORM_TO_MINUS1_1: y = (y.clamp(-1,1) + 1.0)/2.0 else: y = y.clamp(0,1) y = (y*255.0).byte().numpy().transpose(1,2,0) return Image.fromarray(y, mode="RGB") # ---------------- 路由 ---------------- @app.route("/", methods=["GET"]) def home(): return render_template_string(INDEX_HTML, weight_name=os.path.basename(WEIGHTS_PATH)) @app.route("/api/predict", methods=["OPTIONS"]) def predict_preflight(): return ("", 204) import time, uuid, os @app.route("/api/predict", methods=["POST"]) @torch.no_grad() def predict(): app.logger.info("== /api/predict called at %s ==", time.time()) if NETG is None: return jsonify({"ok": False, "error": "model not ready"}), 500 f = request.files.get("image") if not f: return jsonify({"ok": False, "error": "no image"}), 400 fname = secure_filename(f.filename or f"{uuid.uuid4().hex}.png") inpath = os.path.join(UPLOAD_DIR, fname) f.save(inpath) img = Image.open(inpath) x = pil_to_tensor_4ch(img, mask_img=None).unsqueeze(0).to(DEVICE) # [1,4,H,W] t0 = time.time() y = NETG(x) # 期待 [1,3,H,W] latency = int((time.time() - t0)*1000) out_img = tensor_to_pil(y) out_name = f"{uuid.uuid4().hex}.png" out_path = os.path.join(OUTPUT_DIR, out_name) out_img.save(out_path) return jsonify({"ok": True, "latency_ms": latency, "output_url": f"/outputs/{out_name}"}) @app.route("/outputs/") def outputs(name): return send_from_directory(OUTPUT_DIR, name) # ---------------- 主入口 ---------------- if __name__ == "__main__": print(f"Device: {DEVICE}") init_model() app.run(host="127.0.0.1", port=5000, debug=True) @app.errorhandler(403) def handle_403(e): app.logger.warning("403 Forbidden: %s", repr(e)) return jsonify({"ok": False, "error": "forbidden", "detail": str(e)}), 403