Scene stuttering/jitter caused by simGetImages calls

import json
import logging
import os
import signal
import subprocess
import sys
import time
import traceback

import msgpackrpc
import airsim

log = logging.getLogger("stream_worker")

running = True


def stop(*_):
    global running
    running = False


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--spec-file", help="JSON 参数文件路径（网关使用）")
    parser.add_argument("--spec", help="内联 JSON 参数（手动调试用，二选一）")
    return parser.parse_args()


def load_spec(args) -> dict:
    if args.spec_file:
        with open(args.spec_file, "r", encoding="utf-8") as f:
            return json.load(f)
    if args.spec:
        return json.loads(args.spec)
    raise SystemExit("stream_worker: 需要 --spec-file 或 --spec")


def build_ffmpeg(spec: dict) -> subprocess.Popen:
    codec = "h264_nvenc" if spec.get("gpuEncoder") else "libx264"
    fps = int(spec.get("fps", 25))
    image_format = spec.get("imageFormat", "png").lower()
    size = f"{int(spec.get('width', 1280))}x{int(spec.get('height', 720))}"
    bitrate = str(spec.get("bitrate", "2M"))

    cmd = ["ffmpeg", "-loglevel", "warning"]
    if image_format == "rawvideo":
        # 原始像素直接编码（无需解码），省 AirSim PNG 压缩，降低取帧延迟
        cmd += ["-f", "rawvideo", "-pix_fmt", "bgr24",
                "-s", size, "-r", str(fps), "-i", "-",
                "-c:v", codec, "-pix_fmt", "yuv420p", "-b:v", bitrate]
    else:
        # png/jpeg 走 image2pipe 解码
        input_vcodec = "mjpeg" if image_format == "jpeg" else "png"
        cmd += ["-f", "image2pipe", "-vcodec", input_vcodec,
                "-s", size, "-r", str(fps), "-i", "-",
                "-c:v", codec, "-pix_fmt", "yuv420p", "-b:v", bitrate]

    # 低延迟参数：适用于实时推流场景
    if codec == "libx264":
        cmd.extend([
            "-preset", "ultrafast",      # 跳过复杂运动搜索，移动时编码耗时 -60~80%
            "-tune", "zerolatency",      # 消除编码器内部缓冲，帧写完立即输出
            "-bf", "0",                  # 禁用 B 帧，消除无效的时域预测等待
            "-g", str(fps),              # GOP = 1 秒帧数，控制 I 帧间隔
        ])
    elif codec == "h264_nvenc":
        # NVENC 参数（需要 NVIDIA 驱动 >= 470.x）
        cmd.extend([
            "-preset", "p1",             # NVENC 最快预设
            "-tune", "ull",              # Ultra Low Latency
            "-bf", "0",
            "-g", str(fps),
        ])

    cmd.extend(["-f", "flv", spec["rtmpUrl"]])

    return subprocess.Popen(cmd, stdin=subprocess.PIPE)


def connect_airsim(spec: dict) -> airsim.MultirotorClient:
    host = spec.get("airsimHost") or os.getenv("AIRSIM_HOST", "127.0.0.1")
    port = int(spec.get("airsimPort") or os.getenv("AIRSIM_PORT", "41451"))
    timeout = float(spec.get("airsimTimeout") or os.getenv("AIRSIM_TIMEOUT_SEC", "5"))
    log.info("Connecting to AirSim: %s:%s, timeout=%ss", host, port, timeout)
    client = airsim.MultirotorClient(ip=host, port=port, timeout_value=timeout)
    try:
        client.confirmConnection()
        log.info("Connected to AirSim successfully")
    except Exception as exc:
        log.error("FAILED to connect to AirSim: %s: %s", type(exc).__name__, exc)
        log.error("请检查：AirSim 是否运行 / RPC 是否监听 %s:%s / 防火墙是否放行", host, port)
        raise
    return client


def main():
    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    if hasattr(signal, "SIGBREAK"):  # Windows: 父进程通过 CTRL_BREAK_EVENT 触发优雅停止
        signal.signal(signal.SIGBREAK, stop)

    logging.basicConfig(
        stream=sys.stderr,
        level=logging.INFO,
        format="%(asctime)s [stream_worker] %(message)s",
    )

    spec = load_spec(parse_args())
    camera_name = spec["cameraName"]
    vehicle_name = spec.get("vehicleName", "")
    fps = int(spec.get("fps", 25))
    image_format = spec.get("imageFormat", "png").lower()

    client = connect_airsim(spec)
    ffmpeg = build_ffmpeg(spec)

    interval = 1.0 / max(fps, 1)
    # AirSim 启动后相机可能需要几帧初始化，允许几次连续超时重试
    max_retries = 3
    retry_count = 0
    # 只通知一次：Gateway stream_start 在等这条 stdout 事件后再返回成功
    ready_notified = False

    try:
        while running:
            loop_start = time.monotonic()
            try:
                if image_format == "rawvideo":
                    # 原始 RGB 像素（compress=False），省 AirSim PNG 压缩，降低取帧延迟
                    responses = client.simGetImages([
                        airsim.ImageRequest(camera_name, airsim.ImageType.Scene, pixels_as_float=False, compress=False)
                    ], vehicle_name=vehicle_name if vehicle_name else "")
                    raw_image = responses[0].image_data_uint8 if responses else None
                elif image_format == "jpeg":
                    # 使用 simGetImages 获取 JPEG 压缩图像
                    responses = client.simGetImages([
                        airsim.ImageRequest(camera_name, airsim.ImageType.Scene, compress=True)
                    ], vehicle_name=vehicle_name if vehicle_name else "")
                    raw_image = responses[0].image_data_uint8 if responses else None
                else:
                    # 默认 PNG 格式
                    if vehicle_name:
                        raw_image = client.simGetImage(camera_name, airsim.ImageType.Scene, vehicle_name=vehicle_name)
                    else:
                        raw_image = client.simGetImage(camera_name, airsim.ImageType.Scene)

                retry_count = 0
                if raw_image and ffmpeg.stdin:
                    ffmpeg.stdin.write(raw_image)
                    ffmpeg.stdin.flush()
                    # ---------------------------------------------------------------
                    # 通知父进程「已真正开始推流」：仅进程拉起不算成功。
                    # Gateway 收到后才结束 stream.start，避免上游过早 reply、前端 HLS 对着空流空转。
                    # ---------------------------------------------------------------
                    if not ready_notified:
                        print("STREAM_READY", flush=True)
                        ready_notified = True

            except msgpackrpc.error.TimeoutError:
                retry_count += 1
                log.warning("simGetImage timeout (%s/%s), camera=%s, vehicle=%s",
                            retry_count, max_retries, camera_name, vehicle_name)
                if retry_count >= max_retries:
                    log.error("FAILED: %s consecutive timeouts, exiting", max_retries)
                    break
                time.sleep(1)
                continue
            except Exception as exc:
                log.error("FAILED: simGetImage exception: %s: %s", type(exc).__name__, exc)
                traceback.print_exc()
                break

            # 延迟优化第二项 sleep 扣耗时：扣除本轮取帧+写管道已耗时间，避免固定 sleep 导致帧率虚降
            elapsed = time.monotonic() - loop_start
            time.sleep(max(0, interval - elapsed))
    finally:
        if ffmpeg.stdin:
            ffmpeg.stdin.close()
        ffmpeg.terminate()
        try:
            ffmpeg.wait(timeout=5)
        except subprocess.TimeoutExpired:
            ffmpeg.kill()


if __name__ == "__main__":
    main()