NAFNet/python/axmodel_infer.py

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

import os

import cv2
import numpy as np
import axengine as axe

# ============================================================
# 默认参数（按需修改）
# ============================================================
AXMODEL_PATH = 'NAFNet_1_3_256_256.axmodel'
INPUT_PATH = 'demo/noisy.png'
OUTPUT_PATH = 'axmodel_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.uint8)
    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, 'noisy', (int(10 * font_scale), int(30 * font_scale)),
                font, font_scale, color, thickness, cv2.LINE_AA)
    cv2.putText(output_bgr, 'denoised', (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)

    session = axe.InferenceSession(AXMODEL_PATH, providers=['AxEngineExecutionProvider'])
    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 axmodel shape {expected_shape}.')

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


if __name__ == '__main__':
    main()