#!/usr/bin/env python3
"""Single-image NAFNet ONNX inference with side-by-side comparison output."""

import os

import cv2
import numpy as np
import onnxruntime as ort

# ============================================================
# 默认参数（按需修改）
# ============================================================
ONNX_PATH = 'experiments/onnx/NAFNet-SIDD-width64-256x256.onnx'
INPUT_PATH = 'demo/noisy.png'
OUTPUT_PATH = 'demo/denoise_onnx_compare.png'


def read_image(path):
    bgr = cv2.imread(path, cv2.IMREAD_COLOR)
    if bgr is None:
        raise FileNotFoundError(f'Cannot read image: {path}')
    rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
    tensor = np.transpose(rgb, (2, 0, 1))[None, ...]
    return tensor, bgr


def save_comparison(orig_bgr, output_tensor, path):
    """拼接原图与结果图并标注文字，输出图还原到原图大小后再拼接"""
    output = np.squeeze(output_tensor, axis=0)
    output = np.clip(output, 0.0, 1.0)
    output_rgb = np.transpose(output, (1, 2, 0))
    output_bgr = cv2.cvtColor((output_rgb * 255.0).round().astype(np.uint8), cv2.COLOR_RGB2BGR)

    h, w = orig_bgr.shape[:2]
    if output_bgr.shape[:2] != (h, w):
        output_bgr = cv2.resize(output_bgr, (w, h), interpolation=cv2.INTER_LINEAR)

    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = max(h, w) / 512.0
    thickness = max(1, int(font_scale * 2))
    color = (255, 255, 255)

    cv2.putText(orig_bgr, 'input', (int(10 * font_scale), int(30 * font_scale)),
                font, font_scale, color, thickness, cv2.LINE_AA)
    cv2.putText(output_bgr, 'onnx', (int(10 * font_scale), int(30 * font_scale)),
                font, font_scale, color, thickness, cv2.LINE_AA)

    compare = np.concatenate([orig_bgr, output_bgr], axis=1)
    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
    cv2.imwrite(path, compare)


def main():
    inp, orig_bgr = read_image(INPUT_PATH)

    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if ort.get_device() == 'GPU' else ['CPUExecutionProvider']
    session = ort.InferenceSession(ONNX_PATH, providers=providers)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    expected_shape = session.get_inputs()[0].shape
    if list(inp.shape) != expected_shape:
        raise ValueError(f'Input shape {list(inp.shape)} does not match fixed ONNX shape {expected_shape}.')

    out = session.run([output_name], {input_name: inp.astype(np.float32)})[0]
    save_comparison(orig_bgr, out, OUTPUT_PATH)
    print(f'ONNX inference finished: {OUTPUT_PATH}')


if __name__ == '__main__':
    main()