| |
| """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() |
|
|
