| import argparse |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import onnxruntime as ort |
|
|
| def hwc_to_chw(img): |
| return np.transpose(img, axes=[2, 0, 1]).copy() |
|
|
|
|
| def chw_to_hwc(img): |
| return np.transpose(img, axes=[1, 2, 0]).copy() |
|
|
|
|
| def read_rgb_float(image_path, size=None): |
| img = cv2.imread(str(image_path)) |
| if img is None: |
| raise FileNotFoundError(f'Failed to read image: {image_path}') |
| img = img[:, :, ::-1] |
| if size is not None: |
| img = cv2.resize(img, size, interpolation=cv2.INTER_AREA) |
| return img.astype(np.float32) / 255.0 |
|
|
|
|
| def preprocess(image_path, size=None): |
| img = read_rgb_float(image_path, size=size) * 2 - 1 |
| return hwc_to_chw(img).astype(np.float32)[None, ...] |
|
|
|
|
| def get_onnx_input_size(session): |
| shape = session.get_inputs()[0].shape |
| h, w = shape[2], shape[3] |
| if isinstance(h, int) and isinstance(w, int): |
| return w, h |
| return None |
|
|
|
|
| def infer_onnx(session, image_path, size=None): |
| tensor = preprocess(image_path, size=size) |
| input_name = session.get_inputs()[0].name |
| output_name = session.get_outputs()[0].name |
| out = session.run([output_name], {input_name: tensor})[0] |
| out = np.clip(out, -1, 1) |
| out = out * 0.5 + 0.5 |
| return chw_to_hwc(out.squeeze(0)) |
|
|
|
|
| def load_session(onnx_path, use_cpu=False): |
| providers = [] |
| if not use_cpu and 'CUDAExecutionProvider' in ort.get_available_providers(): |
| providers.append('CUDAExecutionProvider') |
| providers.append('CPUExecutionProvider') |
| return ort.InferenceSession(str(onnx_path), providers=providers) |
|
|
|
|
| def save_comparison(hazy_path, dehazed, output_path, size=None): |
| """拼接原图(hazy)和去雾图(dehazed)，并标注标签""" |
| hazy = cv2.imread(str(hazy_path)) |
| if hazy is None: |
| raise FileNotFoundError(f'Failed to read image: {hazy_path}') |
|
|
| h, w = dehazed.shape[:2] |
| if hazy.shape[:2] != (h, w): |
| hazy = cv2.resize(hazy, (w, h), interpolation=cv2.INTER_AREA) |
|
|
| |
| dehazed_bgr = cv2.cvtColor( |
| np.round(dehazed * 255).clip(0, 255).astype(np.uint8), |
| cv2.COLOR_RGB2BGR, |
| ) |
|
|
| pad = 40 |
| canvas = np.full((h + pad, w * 2, 3), 255, dtype=np.uint8) |
| canvas[pad:, :w] = hazy |
| canvas[pad:, w:] = dehazed_bgr |
|
|
| font = cv2.FONT_HERSHEY_SIMPLEX |
| cv2.putText(canvas, 'Hazy', (w // 2 - 40, pad - 8), font, 0.9, (0, 0, 0), 2, cv2.LINE_AA) |
| cv2.putText(canvas, 'Dehazed', (w + w // 2 - 60, pad - 8), font, 0.9, (0, 0, 0), 2, cv2.LINE_AA) |
|
|
| output_path.parent.mkdir(parents=True, exist_ok=True) |
| cv2.imwrite(str(output_path), canvas) |
|
|
|
|
| def main(args): |
| onnx_path = Path(args.onnx) |
| if not onnx_path.is_file(): |
| raise FileNotFoundError(f'ONNX model not found: {onnx_path}') |
|
|
| input_path = Path(args.input) |
| if not input_path.is_file(): |
| raise FileNotFoundError(f'Input image not found: {input_path}') |
|
|
| output_path = Path(args.output) |
| output_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| session = load_session(onnx_path, args.cpu) |
|
|
| onnx_size = get_onnx_input_size(session) |
| if args.width > 0 and args.height > 0: |
| infer_size = (args.width, args.height) |
| elif onnx_size is not None: |
| infer_size = onnx_size |
| else: |
| infer_size = None |
|
|
| print(f'ONNX: {onnx_path}') |
| print(f'Input: {input_path}') |
| print(f'Size: {infer_size}') |
| print(f'Output: {output_path}') |
|
|
| out_img = infer_onnx(session, input_path, size=infer_size) |
| save_comparison(input_path, out_img, output_path) |
| print('Done.') |
|
|
|
|
| if __name__ == '__main__': |
| parser = argparse.ArgumentParser(description='Single-image ONNX inference for DehazeFormer.') |
| parser.add_argument('--input', required=True, type=str, help='path to input image') |
| parser.add_argument('--onnx', required=True, type=str, help='path to ONNX model') |
| parser.add_argument('--output', default='output.png', type=str, help='path to output image') |
| parser.add_argument('--width', default=-1, type=int, help='resize input width (-1 uses ONNX fixed size)') |
| parser.add_argument('--height', default=-1, type=int, help='resize input height (-1 uses ONNX fixed size)') |
| parser.add_argument('--cpu', action='store_true', help='force CPU inference') |
| main(parser.parse_args()) |
|
|
