Spaces:
Sleeping
Sleeping
| # 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 | |
| # 统一给响应加上允许头（有时浏览器仍会要） | |
| 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""" | |
| <!doctype html> | |
| <html lang="zh-CN"> | |
| <head><meta charset="utf-8"/><meta name="viewport" content="width=device-width, initial-scale=1"/> | |
| <title>Pix2Pix 4通道推理 Demo</title> | |
| <style> | |
| body { | |
| font-family: system-ui; | |
| max-width: 980px; | |
| margin: 24px auto; | |
| } | |
| .card { | |
| border: 1px solid #e5e7eb; | |
| border-radius: 12px; | |
| padding: 16px; | |
| margin: 12px 0; | |
| } | |
| .btn { | |
| padding: 10px 16px; | |
| border-radius: 8px; | |
| border: 1px solid #e5e7eb; | |
| background: #10b981; | |
| color: #fff; | |
| cursor: pointer; | |
| } | |
| .mono { | |
| font-family: ui-monospace, Menlo, Consolas, monospace; | |
| } | |
| /* ✅ 新增部分：限制结果图片的显示尺寸 */ | |
| #result img { | |
| max-width: 128px; /* 最大显示宽度为 512px */ | |
| max-height: 600px; /* 超过则等比例缩放 */ | |
| /* 自适应容器宽度 */ | |
| height: auto; | |
| border-radius: 8px; | |
| border: 1px solid #e5e7eb; | |
| display: block; | |
| margin-top: 10px; | |
| } | |
| /* 可选：添加“查看原图”链接样式 */ | |
| .view-full { | |
| display: inline-block; | |
| margin-top: 8px; | |
| color: #2563eb; | |
| font-size: 14px; | |
| } | |
| </style> | |
| </head> | |
| <body> | |
| <h2>Pix2Pix 4通道（RGB+黑mask, 512×1536）在线推理</h2> | |
| <div class="card"> | |
| <p>当前权重：<code id="w"></code></p> | |
| <p>上传一张图片（将被缩放到 512×1536，并在后端叠加全黑 mask 作为第4通道）：</p> | |
| <input id="img" type="file" accept="image/*"/> | |
| <button id="run" class="btn">生成</button> | |
| <span id="s" class="mono"></span> | |
| </div> | |
| <div class="card"> | |
| <h3>输出</h3> | |
| <div id="out"></div> | |
| </div> | |
| <script> | |
| async function postForm(url, formData){ | |
| const r = await fetch(url, {method:"POST", body:formData}); | |
| if(!r.ok) throw new Error("HTTP "+r.status); | |
| return await r.json(); | |
| } | |
| document.getElementById("w").textContent = "{{ weight_name }}"; | |
| document.getElementById("run").onclick = async ()=>{ | |
| const s = document.getElementById("s"); | |
| const out = document.getElementById("out"); | |
| out.innerHTML = ""; | |
| const f = document.getElementById("img").files[0]; | |
| if(!f){ s.textContent = "请先选择图片"; return; } | |
| s.textContent = "推理中..."; | |
| const fd = new FormData(); | |
| fd.append("image", f); | |
| try{ | |
| const j = await postForm("/api/predict", fd); | |
| s.textContent = "耗时 "+j.latency_ms+" ms"; | |
| const img = document.createElement("img"); | |
| img.src = j.output_url+"?t="+Date.now(); | |
| img.style.maxWidth = "100%"; | |
| out.appendChild(img); | |
| }catch(e){ | |
| s.textContent = "失败："+e.message; | |
| } | |
| }; | |
| </script> | |
| </body></html> | |
| """ | |
| # ---------------- 模型加载（关键） ---------------- | |
| 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") | |
| # ---------------- 路由 ---------------- | |
| def home(): | |
| return render_template_string(INDEX_HTML, weight_name=os.path.basename(WEIGHTS_PATH)) | |
| def predict_preflight(): | |
| return ("", 204) | |
| import time, uuid, os | |
| 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}"}) | |
| 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) | |
| def handle_403(e): | |
| app.logger.warning("403 Forbidden: %s", repr(e)) | |
| return jsonify({"ok": False, "error": "forbidden", "detail": str(e)}), 403 |