File size: 2,521 Bytes
1266aec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | #!/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()
|