451 lines
19 KiB
Python
451 lines
19 KiB
Python
#!/usr/bin/env python3
|
||
"""基准测试：19张图片 × 5种发际线，记录每步耗时和显存变化。
|
||
|
||
用法： python3 benchmark_grow.py
|
||
输出： benchmark_results.json + benchmark_report.html
|
||
"""
|
||
import json, os, time, base64, subprocess, re, glob, shutil
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
import urllib.request, urllib.error
|
||
|
||
API_URL = "http://127.0.0.1:8187/api/v1/hair/grow"
|
||
TOKEN = "dev-shared-secret-2026"
|
||
IMG_DIR = "/home/ubuntu/hair/image/girl_img"
|
||
OUT_DIR = "/home/ubuntu/hair/benchmark_out"
|
||
RESULT_JSON = os.path.join(OUT_DIR, "benchmark_results.json")
|
||
REPORT_HTML = os.path.join(OUT_DIR, "benchmark_report.html")
|
||
|
||
HAIRLINES = [
|
||
("1", "ellipse", "椭圆形"),
|
||
("2", "flower", "花瓣形"),
|
||
("3", "heart", "心形"),
|
||
("4", "straight", "直线形"),
|
||
("5", "wave", "波浪形"),
|
||
]
|
||
|
||
os.makedirs(OUT_DIR, exist_ok=True)
|
||
|
||
|
||
# ── 工具函数 ──────────────────────────────────────────────
|
||
|
||
def get_vram():
|
||
"""返回 (used_MB, free_MB)"""
|
||
try:
|
||
out = subprocess.check_output(
|
||
["nvidia-smi", "--query-gpu=memory.used,memory.free",
|
||
"--format=csv,noheader,nounits"], text=True, timeout=5
|
||
).strip()
|
||
used, free = out.split(",")
|
||
return int(used.strip()), int(free.strip())
|
||
except Exception:
|
||
return -1, -1
|
||
|
||
|
||
def get_gpu_procs():
|
||
"""返回各进程显存占用 dict"""
|
||
try:
|
||
out = subprocess.check_output(
|
||
["nvidia-smi", "--query-compute-apps=pid,used_memory",
|
||
"--format=csv,noheader,nounits"], text=True, timeout=5
|
||
).strip()
|
||
procs = {}
|
||
for line in out.splitlines():
|
||
parts = line.split(",")
|
||
if len(parts) >= 2:
|
||
procs[parts[0].strip()] = int(parts[1].strip())
|
||
return procs
|
||
except Exception:
|
||
return {}
|
||
|
||
|
||
def read_worker_log_tail(n=50):
|
||
"""读取 worker.log 最后 n 行"""
|
||
log_path = "/home/ubuntu/hair/worker.log"
|
||
try:
|
||
result = subprocess.run(["tail", "-n", str(n), log_path],
|
||
capture_output=True, text=True, timeout=5)
|
||
return result.stdout
|
||
except Exception:
|
||
return ""
|
||
|
||
|
||
def call_api(image_path, hair_style_value):
|
||
"""调用接口2，返回 (json_dict, elapsed_sec, error_or_None)"""
|
||
import mimetypes
|
||
boundary = "----BenchmarkBoundary" + str(int(time.time()*1000))
|
||
filename = os.path.basename(image_path)
|
||
mime = mimetypes.guess_type(image_path)[0] or "image/jpeg"
|
||
|
||
with open(image_path, "rb") as f:
|
||
img_data = f.read()
|
||
|
||
body_parts = []
|
||
body_parts.append(f"--{boundary}\r\n".encode())
|
||
body_parts.append(f'Content-Disposition: form-data; name="image_file"; filename="{filename}"\r\n'.encode())
|
||
body_parts.append(f"Content-Type: {mime}\r\n\r\n".encode())
|
||
body_parts.append(img_data)
|
||
body_parts.append(f"\r\n--{boundary}\r\n".encode())
|
||
body_parts.append(b'Content-Disposition: form-data; name="gender"\r\n\r\n')
|
||
body_parts.append(b"female")
|
||
body_parts.append(f"\r\n--{boundary}\r\n".encode())
|
||
body_parts.append(b'Content-Disposition: form-data; name="hair_style"\r\n\r\n')
|
||
body_parts.append(hair_style_value.encode())
|
||
body_parts.append(f"\r\n--{boundary}\r\n".encode())
|
||
body_parts.append(b'Content-Disposition: form-data; name="use_mask"\r\n\r\n')
|
||
body_parts.append(b"0")
|
||
body_parts.append(f"\r\n--{boundary}\r\n".encode())
|
||
body_parts.append(b'Content-Disposition: form-data; name="prompt"\r\n\r\n')
|
||
body_parts.append(b"\xe5\xa1\xab\xe5\x85\x85\xe9\x81\xae\xe7\xbd\xa9\xe5\x8c\xba\xe5\x9f\x9f\xe7\x9a\x84\xe5\xa4\xb4\xe5\x8f\x91\xef\xbc\x8c\xe7\x9a\xae\xe8\x82\xa4\xe5\x8a\xa0\xe4\xb8\x80\xe7\x82\xb9\xe7\xa3\xa8\xe7\x9a\xae")
|
||
body_parts.append(f"\r\n--{boundary}--\r\n".encode())
|
||
body = b"".join(body_parts)
|
||
|
||
req = urllib.request.Request(API_URL, data=body, method="POST")
|
||
req.add_header("Content-Type", f"multipart/form-data; boundary={boundary}")
|
||
req.add_header("X-Internal-Token", TOKEN)
|
||
|
||
t0 = time.time()
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=600) as resp:
|
||
raw = resp.read()
|
||
elapsed = time.time() - t0
|
||
data = json.loads(raw)
|
||
return data, elapsed, None
|
||
except urllib.error.HTTPError as e:
|
||
elapsed = time.time() - t0
|
||
try:
|
||
data = json.loads(e.read())
|
||
except Exception:
|
||
data = {"error": str(e)}
|
||
return data, elapsed, f"HTTP {e.code}"
|
||
except Exception as e:
|
||
elapsed = time.time() - t0
|
||
return {}, elapsed, str(e)
|
||
|
||
|
||
def parse_step_timing(log_text):
|
||
"""从 worker.log 文本中解析步骤耗时"""
|
||
timing = {}
|
||
for line in log_text.splitlines():
|
||
# 匹配 "步骤1 遮罩完成 耗时=949ms"
|
||
m = re.search(r"步骤(\d+)\s+\S+\s+耗时=(\d+)ms", line)
|
||
if m:
|
||
timing[f"step{m.group(1)}_ms"] = int(m.group(2))
|
||
# 匹配 "人脸检出 px_per_cm=30.669 图尺寸=813x967"
|
||
if "人脸检出" in line:
|
||
m2 = re.search(r"px_per_cm=([\d.]+)", line)
|
||
if m2:
|
||
timing["px_per_cm"] = float(m2.group(1))
|
||
# 匹配 "头发分割完成"
|
||
if "头发分割完成" in line:
|
||
timing["hair_seg"] = True
|
||
# 匹配 "换发型图失败"
|
||
if "换发型图失败" in line or "换发型服务不可达" in line:
|
||
timing["swap_error"] = line.strip()[-100:]
|
||
# 匹配 "重绘" 相关
|
||
if "重绘" in line and ("完成" in line or "失败" in line):
|
||
timing["redraw_status"] = "完成" if "完成" in line else "失败"
|
||
return timing
|
||
|
||
|
||
# ── 主流程 ────────────────────────────────────────────────
|
||
|
||
def run_benchmark():
|
||
images = sorted(glob.glob(os.path.join(IMG_DIR, "girl*.jpg")))
|
||
print(f"找到 {len(images)} 张图片")
|
||
|
||
# 加载已有结果（支持断点续跑）
|
||
results = []
|
||
if os.path.exists(RESULT_JSON):
|
||
with open(RESULT_JSON) as f:
|
||
results = json.load(f)
|
||
print(f"已有 {len(results)} 条记录，继续未完成的测试")
|
||
|
||
total_calls = len(images) * len(HAIRLINES)
|
||
done = len(results)
|
||
print(f"总计 {total_calls} 次调用，已完成 {done}，剩余 {total_calls - done}")
|
||
|
||
log_offset = 0
|
||
try:
|
||
log_offset = subprocess.check_output(["wc", "-l", "/home/ubuntu/hair/worker.log"],
|
||
text=True, timeout=5).split()[0]
|
||
log_offset = int(log_offset)
|
||
except Exception:
|
||
pass
|
||
|
||
for img_idx, img_path in enumerate(images):
|
||
img_name = os.path.basename(img_path)
|
||
|
||
# 跳过已完成的图片
|
||
img_results = [r for r in results if r["image"] == img_name]
|
||
if len(img_results) >= len(HAIRLINES):
|
||
print(f"[{img_idx+1}/{len(images)}] {img_name} 已完成，跳过")
|
||
continue
|
||
|
||
for hl_id, hl_key, hl_label in HAIRLINES:
|
||
# 跳过已完成的
|
||
existing = [r for r in results if r["image"] == img_name and r["hairline_id"] == hl_id]
|
||
if existing:
|
||
continue
|
||
|
||
print(f"\n[{img_idx+1}/{len(images)}] {img_name} → {hl_label}({hl_id}) ...")
|
||
|
||
# 记录 worker.log 行数
|
||
try:
|
||
log_before = int(subprocess.check_output(
|
||
["wc", "-l", "/home/ubuntu/hair/worker.log"], text=True, timeout=5).split()[0])
|
||
except Exception:
|
||
log_before = 0
|
||
|
||
# VRAM before
|
||
vram_before, vram_free_before = get_vram()
|
||
procs_before = get_gpu_procs()
|
||
t_start = time.time()
|
||
|
||
# 调用API
|
||
api_result, elapsed, error = call_api(img_path, hl_id)
|
||
|
||
t_end = time.time()
|
||
|
||
# VRAM after
|
||
vram_after, vram_free_after = get_vram()
|
||
procs_after = get_gpu_procs()
|
||
|
||
# 读取新日志
|
||
try:
|
||
log_diff = subprocess.check_output(
|
||
["tail", "-n", "+{}".format(log_before + 1), "/home/ubuntu/hair/worker.log"],
|
||
text=True, timeout=5)
|
||
except Exception:
|
||
log_diff = ""
|
||
|
||
step_timing = parse_step_timing(log_diff)
|
||
|
||
# 提取结果图片
|
||
preview_b64 = ""
|
||
grown_b64 = ""
|
||
api_code = api_result.get("code", -1)
|
||
api_msg = api_result.get("message", "")
|
||
api_results = api_result.get("data", {}).get("results", [])
|
||
if api_results:
|
||
r0 = api_results[0]
|
||
preview_b64 = r0.get("image_base64", "")
|
||
grown_b64 = r0.get("grown_image_base64", "")
|
||
|
||
# 保存缩略图
|
||
thumb_dir = os.path.join(OUT_DIR, "thumbs")
|
||
os.makedirs(thumb_dir, exist_ok=True)
|
||
if grown_b64:
|
||
grown_bytes = base64.b64decode(grown_b64)
|
||
thumb_path = os.path.join(thumb_dir, f"{img_name}_hl{hl_id}_grown.jpg")
|
||
with open(thumb_path, "wb") as f:
|
||
f.write(grown_bytes)
|
||
if preview_b64:
|
||
preview_bytes = base64.b64decode(preview_b64)
|
||
thumb_path = os.path.join(thumb_dir, f"{img_name}_hl{hl_id}_preview.png")
|
||
with open(thumb_path, "wb") as f:
|
||
f.write(preview_bytes)
|
||
|
||
record = {
|
||
"image": img_name,
|
||
"image_idx": img_idx + 1,
|
||
"hairline_id": hl_id,
|
||
"hairline_key": hl_key,
|
||
"hairline_label": hl_label,
|
||
"timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||
"total_time_s": round(elapsed, 2),
|
||
"api_code": api_code,
|
||
"api_msg": api_msg,
|
||
"error": error,
|
||
"has_preview": bool(preview_b64),
|
||
"has_grown": bool(grown_b64),
|
||
"preview_size": len(preview_b64),
|
||
"grown_size": len(grown_b64),
|
||
"vram_before_mb": vram_before,
|
||
"vram_after_mb": vram_after,
|
||
"vram_free_before_mb": vram_free_before,
|
||
"vram_free_after_mb": vram_free_after,
|
||
"vram_delta_mb": vram_after - vram_before,
|
||
"procs_before": procs_before,
|
||
"procs_after": procs_after,
|
||
"step_timing": step_timing,
|
||
}
|
||
results.append(record)
|
||
print(f" → HTTP code={api_code} time={elapsed:.1f}s vram={vram_before}→{vram_after}MB "
|
||
f"preview={'✓' if preview_b64 else '✗'} grown={'✓' if grown_b64 else '✗'}")
|
||
|
||
# 保存中间结果
|
||
with open(RESULT_JSON, "w") as f:
|
||
json.dump(results, f, indent=2, ensure_ascii=False)
|
||
|
||
# 生成中间报告
|
||
generate_html(results)
|
||
|
||
# 最终报告
|
||
generate_html(results)
|
||
print(f"\n✅ 完成！结果: {RESULT_JSON}")
|
||
print(f"📄 报告: {REPORT_HTML}")
|
||
print(f"📊 总调用: {len(results)}/{total_calls}")
|
||
|
||
|
||
def generate_html(results):
|
||
"""生成HTML报告"""
|
||
# 统计
|
||
total = len(results)
|
||
success = sum(1 for r in results if r["has_grown"])
|
||
failed = total - success
|
||
times = [r["total_time_s"] for r in results if r["has_grown"]]
|
||
avg_time = sum(times) / len(times) if times else 0
|
||
max_time = max(times) if times else 0
|
||
min_time = min(times) if times else 0
|
||
|
||
# 按发际线类型分组统计
|
||
hl_stats = {}
|
||
for r in results:
|
||
if r["has_grown"]:
|
||
hl = r["hairline_label"]
|
||
if hl not in hl_stats:
|
||
hl_stats[hl] = {"count": 0, "times": []}
|
||
hl_stats[hl]["count"] += 1
|
||
hl_stats[hl]["times"].append(r["total_time_s"])
|
||
|
||
# 按图片分组
|
||
img_groups = {}
|
||
for r in results:
|
||
img = r["image"]
|
||
if img not in img_groups:
|
||
img_groups[img] = []
|
||
img_groups[img].append(r)
|
||
|
||
# 生成VRAM变化数据
|
||
vram_data = [(i, r["vram_after_mb"]) for i, r in enumerate(results)]
|
||
|
||
html = f"""<!DOCTYPE html>
|
||
<html lang="zh-CN">
|
||
<head>
|
||
<meta charset="UTF-8">
|
||
<title>接口2 基准测试报告</title>
|
||
<style>
|
||
* {{ margin:0; padding:0; box-sizing:border-box; }}
|
||
body {{ font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',Roboto,sans-serif; background:#1a1a2e; color:#e0e0e0; padding:20px; }}
|
||
h1 {{ text-align:center; margin-bottom:20px; color:#00d4ff; }}
|
||
.subtitle {{ text-align:center; color:#888; margin-bottom:30px; font-size:14px; }}
|
||
.summary {{ display:grid; grid-template-columns:repeat(auto-fit,minmax(180px,1fr)); gap:15px; margin-bottom:30px; }}
|
||
.card {{ background:#16213e; border-radius:12px; padding:20px; text-align:center; border:1px solid #333; }}
|
||
.card .num {{ font-size:32px; font-weight:700; }}
|
||
.card .label {{ font-size:12px; color:#888; margin-top:5px; }}
|
||
.card.ok .num {{ color:#0f0; }}
|
||
.card.err .num {{ color:#f44; }}
|
||
.card.time .num {{ color:#00d4ff; }}
|
||
table {{ width:100%; border-collapse:collapse; margin-bottom:30px; background:#16213e; border-radius:12px; overflow:hidden; }}
|
||
th {{ background:#0f3460; padding:12px 8px; text-align:center; font-size:13px; color:#fff; }}
|
||
td {{ padding:8px; text-align:center; border-bottom:1px solid #222; font-size:13px; }}
|
||
tr:hover {{ background:#1a1a3e; }}
|
||
.img-cell {{ text-align:left; }}
|
||
.time-bar {{ display:inline-block; height:20px; background:linear-gradient(90deg,#0f3460,#00d4ff); border-radius:4px; vertical-align:middle; min-width:2px; }}
|
||
.hl-badge {{ display:inline-block; padding:2px 8px; border-radius:10px; font-size:11px; font-weight:600; }}
|
||
.hl-ellipse {{ background:#1b4332; color:#52b788; }}
|
||
.hl-flower {{ background:#3a0ca3; color:#c77dff; }}
|
||
.hl-heart {{ background:#6a040f; color:#ff6b6b; }}
|
||
.hl-straight {{ background:#0077b6; color:#90e0ef; }}
|
||
.hl-wave {{ background:#9d4edd; color:#e0aaff; }}
|
||
.ok {{ color:#0f0; }} .fail {{ color:#f44; }}
|
||
.vram-chart {{ margin:20px 0; }}
|
||
.vram-bars {{ display:flex; align-items:flex-end; height:120px; gap:2px; padding:10px; background:#0d1117; border-radius:8px; }}
|
||
.vram-bar {{ flex:1; background:linear-gradient(180deg,#00d4ff,#0f3460); border-radius:2px 2px 0 0; min-height:2px; position:relative; }}
|
||
.vram-bar:hover::after {{ content:attr(data-val) 'MB'; position:absolute; bottom:100%; left:50%; transform:translateX(-50%); background:#333; padding:2px 6px; border-radius:4px; font-size:10px; white-space:nowrap; }}
|
||
.section-title {{ font-size:18px; font-weight:600; margin:30px 0 15px; color:#00d4ff; border-bottom:1px solid #333; padding-bottom:10px; }}
|
||
.hl-stats {{ display:grid; grid-template-columns:repeat(5,1fr); gap:15px; margin-bottom:20px; }}
|
||
.hl-card {{ background:#16213e; border-radius:12px; padding:15px; text-align:center; }}
|
||
.hl-card .avg {{ font-size:24px; font-weight:700; color:#00d4ff; }}
|
||
.hl-card .minmax {{ font-size:11px; color:#888; margin-top:5px; }}
|
||
.thumb {{ max-width:100px; max-height:100px; border-radius:4px; cursor:pointer; }}
|
||
.thumb:hover {{ transform:scale(2); transition:transform 0.3s; }}
|
||
</style>
|
||
</head>
|
||
<body>
|
||
<h1>📊 接口2 基准测试报告</h1>
|
||
<div class="subtitle">生成时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')} ｜ {total} 次调用</div>
|
||
|
||
<div class="summary">
|
||
<div class="card ok"><div class="num">{success}</div><div class="label">成功</div></div>
|
||
<div class="card err"><div class="num">{failed}</div><div class="label">失败</div></div>
|
||
<div class="card time"><div class="num">{avg_time:.1f}s</div><div class="label">平均耗时</div></div>
|
||
<div class="card time"><div class="num">{min_time:.1f}s</div><div class="label">最快</div></div>
|
||
<div class="card time"><div class="num">{max_time:.1f}s</div><div class="label">最慢</div></div>
|
||
</div>
|
||
"""
|
||
|
||
# 按发际线类型统计
|
||
if hl_stats:
|
||
html += '<div class="section-title">按发际线类型统计</div><div class="hl-stats">'
|
||
for hl_label in ["椭圆形","花瓣形","心形","直线形","波浪形"]:
|
||
if hl_label in hl_stats:
|
||
s = hl_stats[hl_label]
|
||
times = s["times"]
|
||
avg = sum(times) / len(times)
|
||
mn, mx = min(times), max(times)
|
||
cls = hl_label
|
||
html += f'<div class="hl-card"><div class="hl-badge hl-{{cls}}">{hl_label}</div><div class="avg">{avg:.1f}s</div><div class="minmax">{mn:.1f}~{mx:.1f}s ({s["count"]}次)</div></div>'
|
||
else:
|
||
html += f'<div class="hl-card"><div class="hl-badge">{hl_label}</div><div class="avg">-</div><div class="minmax">未完成</div></div>'
|
||
html += '</div>'
|
||
|
||
# VRAM 变化图
|
||
if vram_data:
|
||
max_vram = max(v for _, v in vram_data if v > 0) or 1
|
||
html += '<div class="section-title">显存变化</div><div class="vram-chart"><div class="vram-bars">'
|
||
for i, (_, vram) in enumerate(vram_data):
|
||
if vram > 0:
|
||
h = int(vram / max_vram * 100)
|
||
html += f'<div class="vram-bar" style="height:{h}%" data-val="{vram}" title="第{i+1}次"></div>'
|
||
else:
|
||
html += f'<div class="vram-bar" style="height:0%" data-val="0"></div>'
|
||
html += '</div></div>'
|
||
|
||
# 详细结果表格
|
||
html += '<div class="section-title">详细结果</div><table><thead><tr>'
|
||
html += '<th>#</th><th>图片</th><th>发际线</th><th>总耗时</th><th>遮罩步骤</th>'
|
||
html += '<th>显存前</th><th>显存后</th><th>显存变化</th>'
|
||
html += '<th>预览图</th><th>生发图</th><th>状态</th>'
|
||
html += '</tr></thead><tbody>'
|
||
|
||
for i, r in enumerate(results):
|
||
hl_cls = r["hairline_key"]
|
||
status = '<span class="ok">✓ 成功</span>' if r["has_grown"] else f'<span class="fail">✗ {r.get("error","")}</span>'
|
||
step_ms = r.get("step_timing", {}).get("step1_ms", "")
|
||
step_str = f"{step_ms}ms" if step_ms else "-"
|
||
vram_d = r["vram_delta_mb"]
|
||
vram_d_str = f'<span style="color:{"#f44" if vram_d>0 else "#0f0"}">{"+" if vram_d>=0 else ""}{vram_d}</span>'
|
||
|
||
# 缩略图
|
||
thumb_grown = ""
|
||
if r["has_grown"]:
|
||
thumb_path = f"thumbs/{r['image']}_hl{r['hairline_id']}_grown.jpg"
|
||
if os.path.exists(os.path.join(OUT_DIR, thumb_path)):
|
||
thumb_grown = f'<img class="thumb" src="{thumb_path}">'
|
||
|
||
html += f"""<tr>
|
||
<td>{i+1}</td>
|
||
<td class="img-cell">{r['image']}</td>
|
||
<td><span class="hl-badge hl-{hl_cls}">{r['hairline_label']}</span></td>
|
||
<td><div class="time-bar" style="width:{min(r['total_time_s'],300)}px"></div> {r['total_time_s']:.1f}s</td>
|
||
<td>{step_str}</td>
|
||
<td>{r['vram_before_mb']}MB</td>
|
||
<td>{r['vram_after_mb']}MB</td>
|
||
<td>{vram_d_str}</td>
|
||
<td>{'✓' if r['has_preview'] else '✗'}</td>
|
||
<td>{thumb_grown if thumb_grown else ('✓' if r['has_grown'] else '✗')}</td>
|
||
<td>{status}</td>
|
||
</tr>"""
|
||
|
||
html += '</tbody></table>'
|
||
html += '</body></html>'
|
||
|
||
with open(REPORT_HTML, "w", encoding="utf-8") as f:
|
||
f.write(html)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
run_benchmark()
|