ykxmz/_tests/bench_agent.py

# -*- coding: utf-8 -*-
"""量一下被控端的实际性能，别凭感觉优化。

要测：
    1. 截屏一帧要多久（不同质量/宽度）
    2. 编码一帧要多久
    3. 空闲时 CPU / 内存
    4. 各种质量档位下的字节数（影响传输耗时）
"""
from __future__ import annotations

import io
import os
import statistics
import sys
import time
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT / "windows"))

import psutil  # noqa: E402


def bench_capture(quality: int, width: int, rounds: int = 12):
    from PIL import ImageGrab
    times = []
    sizes = []
    for _ in range(rounds):
        start = time.perf_counter()
        shot = ImageGrab.grab()
        if width and shot.width > width:
            ratio = width / shot.width
            shot = shot.resize((width, int(shot.height * ratio)),
                               ImageGrab.grab().resize.__self__.__class__ and 1
                               or 1)
        buffer = io.BytesIO()
        shot.convert("RGB").save(buffer, format="JPEG", quality=quality)
        times.append((time.perf_counter() - start) * 1000)
        sizes.append(len(buffer.getvalue()))
    return times, sizes


def bench_resize_and_encode(quality: int, width: int, rounds: int = 12):
    """先截一次，反复测 缩放+编码 的耗时（隔离出真正的开销）。"""
    from PIL import Image, ImageGrab
    base = ImageGrab.grab()
    times = []
    sizes = []
    for _ in range(rounds):
        start = time.perf_counter()
        shot = base
        if width and base.width > width:
            ratio = width / base.width
            shot = base.resize((width, int(base.height * ratio)), Image.LANCZOS)
        buffer = io.BytesIO()
        shot.convert("RGB").save(buffer, format="JPEG", quality=quality)
        times.append((time.perf_counter() - start) * 1000)
        sizes.append(len(buffer.getvalue()))
    return times, sizes


def main() -> int:
    from PIL import Image

    screen = Image.new("RGB", (1, 1))
    try:
        from PIL import ImageGrab
        base = ImageGrab.grab()
        screen = base
    except Exception as exc:                              # noqa: BLE001
        print(f"  截屏失败: {exc}")
        return 1

    print(f"  屏幕: {screen.width}x{screen.height}")
    proc = psutil.Process(os.getpid())
    print(f"  基准进程内存: {proc.memory_info().rss / 1024 / 1024:.1f} MB")
    print()

    print("=" * 66)
    print("  一、纯截屏耗时（ImageGrab）")
    print("=" * 66)
    times = []
    for _ in range(15):
        start = time.perf_counter()
        ImageGrab.grab()
        times.append((time.perf_counter() - start) * 1000)
    times.sort()
    print(f"  平均 {statistics.mean(times):6.1f} ms   中位 {statistics.median(times):6.1f} ms   "
          f"最快 {times[0]:6.1f}   最慢 {times[-1]:6.1f}")
    print(f"  -> 理论上限约 {1000 / statistics.median(times):.1f} 帧/秒")

    print()
    print("=" * 66)
    print("  二、缩放 + JPEG 编码耗时（截屏不算）")
    print("=" * 66)
    print(f"  {'宽度':>6} {'质量':>5} {'中位耗时':>10} {'大小':>10} {'理论帧率':>10}")
    for width in (0, 1280, 960, 800):
        for quality in (75, 55, 40):
            t, s = bench_resize_and_encode(quality, width)
            median = statistics.median(t)
            print(f"  {str(width or '原始'):>6} {quality:>5} {median:>9.1f}ms "
                  f"{statistics.mean(s) / 1024:>9.1f}KB {1000 / median:>9.1f}")

    print()
    print("=" * 66)
    print("  三、整条链路（截屏 + 缩放 + 编码）")
    print("=" * 66)
    for width, quality in ((0, 75), (1280, 75), (960, 60), (800, 50)):
        t, s = bench_capture(quality, width)
        median = statistics.median(t)
        print(f"  宽度 {str(width or '原始'):>5} 质量 {quality:>3}: "
              f"{median:>7.1f}ms/帧 -> {1000 / median:>5.1f} 帧/秒, "
              f"{statistics.mean(s) / 1024:.0f}KB/帧, "
              f"{statistics.mean(s) / 1024 * 1000 / median * 8 / 1000:.1f} Mbps")

    print()
    print("=" * 66)
    print("  四、当前默认设置是多少")
    print("=" * 66)
    from lanremote.config import Store
    import tempfile
    store = Store(Path(tempfile.mkdtemp()))
    for key in ("image_width", "image_quality", "screen_fps",
                "screen_quality", "screen_width", "screen_scale"):
        value = store.get(key)
        if value is not None:
            print(f"  {key} = {value}")
    return 0


if __name__ == "__main__":
    sys.exit(main())