fastDVDnet/python/onnx_video_infer.py

#!/usr/bin/env python
"""
FastDVDnet ONNX 视频推理脚本，输出 原帧 | ONNX 去噪结果 拼接视频。

输入视频逐帧 resize 到 ONNX 固定输入尺寸 (HxW)，组成 5 帧窗口
[t-2,t-1,t,t+1,t+2] 推理，输出 resize 回原始尺寸生成拼接视频。

用法:
    python onnx_video_infer.py \
        --onnx ./fastdvdnet_640x480.onnx \
        --video mp4/drone.mp4 \
        --noise_sigma 25 \
        --out_dir ./video_infer_results/onnx_test
"""
import argparse
import os
import time

import cv2
import numpy as np
import onnxruntime as ort


def get_onnx_shapes(onnx_path):
    import onnx
    m = onnx.load(onnx_path)
    x_dims = [d.dim_value for d in m.graph.input[0].type.tensor_type.shape.dim]
    return x_dims  # [N, 15, H, W]


def read_video(video_path, max_frames=0):
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        raise RuntimeError("failed to open video: {}".format(video_path))
    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    orig_w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    orig_h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    frames_bgr = []
    while True:
        ok, frame = cap.read()
        if not ok:
            break
        frames_bgr.append(frame)
        if max_frames and len(frames_bgr) >= max_frames:
            break
    cap.release()
    if not frames_bgr:
        raise RuntimeError("no frames read from {}".format(video_path))
    return frames_bgr, fps, orig_w, orig_h


def reflect_index(idx, length):
    if length <= 1:
        return 0
    while idx < 0 or idx >= length:
        if idx < 0:
            idx = -idx
        if idx >= length:
            idx = 2 * (length - 1) - idx
    return idx


def bgr_to_onnx_input(frame_bgr, onnx_h, onnx_w):
    """resize BGR uint8 to ONNX RGB float32 [1,3,H,W] in [0,1]"""
    resized = cv2.resize(frame_bgr, (onnx_w, onnx_h), interpolation=cv2.INTER_AREA)
    rgb = cv2.cvtColor(resized, cv2.COLOR_BGR2RGB)
    chw = rgb.astype(np.float32).transpose(2, 0, 1) / 255.0
    return chw  # [3, H, W]


def chw_to_bgr_uint8(chw, target_w, target_h):
    """[3,H,W] float32 in [0,1] -> BGR uint8 resized to target_w x target_h"""
    hwc = (chw * 255.0).clip(0, 255).astype(np.uint8).transpose(1, 2, 0)
    rgb = hwc  # already RGB from ONNX output
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    if bgr.shape[1] != target_w or bgr.shape[0] != target_h:
        bgr = cv2.resize(bgr, (target_w, target_h), interpolation=cv2.INTER_LINEAR)
    return bgr


def denoise_onnx(frames_bgr, noise_sigma_01, sess, x_shape, input_names, output_names):
    numframes = len(frames_bgr)
    _, _, onnx_h, onnx_w = x_shape
    temp_psz, ctrl = 5, 2

    # 缓存 resize 后的 CHW 帧
    chw_cache = {}
    def get_chw(i):
        i = i % numframes
        if i not in chw_cache:
            chw_cache[i] = bgr_to_onnx_input(frames_bgr[reflect_index(i, numframes)], onnx_h, onnx_w)
        return chw_cache[i]

    den_frames_bgr = []
    inframes = []  # numpy [3,H,W] each

    for fridx in range(numframes):
        if not inframes:
            for off in range(temp_psz):
                inframes.append(get_chw(fridx + off - ctrl))
        else:
            del inframes[0]
            inframes.append(get_chw(fridx + ctrl))

        # concat 5 frames -> [1, 15, H, W]
        noisy = np.concatenate(inframes, axis=0)[None, :, :, :].astype(np.float32)
        noise_map = np.full((1, 1, onnx_h, onnx_w), noise_sigma_01, dtype=np.float32)
        feeds = {input_names[0]: noisy, input_names[1]: noise_map}
        out = sess.run(output_names, feeds)[0]  # [1, 3, H, W]
        out = np.clip(out, 0.0, 1.0)
        den_bgr = chw_to_bgr_uint8(out[0], frames_bgr[0].shape[1], frames_bgr[0].shape[0])
        den_frames_bgr.append(den_bgr)

    return den_frames_bgr


def write_side_by_side(frames_orig, frames_den, out_path, fps, label=True):
    h, w = frames_orig[0].shape[:2]
    os.makedirs(os.path.dirname(os.path.abspath(out_path)), exist_ok=True)
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    writer = cv2.VideoWriter(out_path, fourcc, fps, (w * 2, h))
    if not writer.isOpened():
        raise RuntimeError("failed to create video writer: {}".format(out_path))
    for orig, den in zip(frames_orig, frames_den):
        canvas = np.concatenate([orig, den], axis=1)
        if label:
            cv2.putText(canvas, "Original", (16, 34), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2, cv2.LINE_AA)
            cv2.putText(canvas, "ONNX Denoised", (w + 16, 34), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 255, 255), 2, cv2.LINE_AA)
        writer.write(canvas)
    writer.release()


def main():
    parser = argparse.ArgumentParser(description="FastDVDnet ONNX video inference side-by-side")
    parser.add_argument("--onnx", type=str, default="./fastdvdnet_640x480.onnx")
    parser.add_argument("--video", type=str, required=True, help="input mp4 video")
    parser.add_argument("--noise_sigma", type=float, default=25.0, help="noise sigma 0-255")
    parser.add_argument("--out_dir", type=str, default="./video_infer_results/onnx_test")
    parser.add_argument("--max_frames", type=int, default=0, help="0 means all frames")
    parser.add_argument("--no_label", action="store_true")
    args = parser.parse_args()

    x_shape = get_onnx_shapes(args.onnx)
    _, _, onnx_h, onnx_w = x_shape
    print("ONNX input shape: {}".format(x_shape))

    sess = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
    input_names = [i.name for i in sess.get_inputs()]
    output_names = [o.name for o in sess.get_outputs()]
    print("providers:", sess.get_providers())
    print("inputs:", [(n, list(i.shape)) for n, i in zip(input_names, sess.get_inputs())])
    print("outputs:", [(n, list(o.shape)) for n, o in zip(output_names, sess.get_outputs())])

    base = os.path.splitext(os.path.basename(args.video))[0]
    out_path = os.path.join(args.out_dir, "{}_onnx_side_by_side.mp4".format(base))

    t0 = time.time()
    frames_bgr, fps, orig_w, orig_h = read_video(args.video, args.max_frames)
    print("video: {}x{} @ {:.1f}fps, {} frames".format(orig_w, orig_h, fps, len(frames_bgr)))
    print("onnx input: {}x{}".format(onnx_w, onnx_h))

    sigma01 = args.noise_sigma / 255.0
    den = denoise_onnx(frames_bgr, sigma01, sess, x_shape, input_names, output_names)
    write_side_by_side(frames_bgr, den, out_path, fps, label=not args.no_label)
    dt = time.time() - t0

    print("[OK] {} -> {}".format(args.video, out_path))
    print("     frames={} time={:.2f}s fps={:.2f}".format(len(frames_bgr), dt, len(frames_bgr) / max(dt, 0.001)))


if __name__ == "__main__":
    main()