File size: 4,357 Bytes
4dca198 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | 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] # BGR -> RGB
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 # [N, C, H, W]
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)
# hazy 是 BGR，dehazed 是 RGB -> 转 BGR 再拼接
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())
|