File size: 3,135 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 | import argparse
import os
import cv2
import numpy as np
def read_image(path, height, width):
image = cv2.imread(path, cv2.IMREAD_COLOR)
if image is None:
raise FileNotFoundError(path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
if image.shape[:2] != (height, width):
image = cv2.resize(image, (width, height), interpolation=cv2.INTER_LINEAR)
image = image * 2.0 - 1.0
image = np.transpose(image, (2, 0, 1))[None].copy()
return image
def save_comparison(hazy_orig, output, path):
"""拼接原图(hazy)与去雾图(dehazed)并标注文字，还原到原图大小"""
orig_h, orig_w = hazy_orig.shape[:2]
# output: NCHW float32 [-1, 1] → 还原到原图大小 HWC uint8 BGR
output = np.clip(output[0], -1.0, 1.0)
output = output * 0.5 + 0.5
output = (np.transpose(output, (1, 2, 0)) * 255).astype(np.uint8)
if output.shape[:2] != (orig_h, orig_w):
output = cv2.resize(output, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)
h, w = hazy_orig.shape[:2]
font = cv2.FONT_HERSHEY_SIMPLEX
font_scale = max(h, w) / 512.0
thickness = max(1, int(font_scale * 2))
color = (255, 255, 255)
# 加 hazy 标注
cv2.putText(hazy_orig, 'hazy', (int(10 * font_scale), int(30 * font_scale)),
font, font_scale, color, thickness, cv2.LINE_AA)
# 加 dehazed 标注
cv2.putText(output, 'dehazed', (int(10 * font_scale), int(30 * font_scale)),
font, font_scale, color, thickness, cv2.LINE_AA)
compare = np.concatenate([hazy_orig, output], axis=1)
os.makedirs(os.path.dirname(path) or '.', exist_ok=True)
cv2.imwrite(path, compare)
def main():
parser = argparse.ArgumentParser(description='Run MixDehazeNet ONNX inference.')
parser.add_argument('--onnx', default='onnx/MixDehazeNet-s_256x256.onnx', help='Path to simplified ONNX model.')
parser.add_argument('--image', default='0003_0.8_0.2.jpg', help='Input hazy image.')
parser.add_argument('--height', type=int, default=256)
parser.add_argument('--width', type=int, default=256)
parser.add_argument('--output', default='./onnx_res.png')
args = parser.parse_args()
try:
import onnxruntime as ort
except ImportError as exc:
raise SystemExit('Please install ONNX Runtime first: pip install onnxruntime-gpu or onnxruntime') from exc
# 读取原图用于显示（BGR uint8）
hazy_orig = cv2.imread(args.image, cv2.IMREAD_COLOR)
if hazy_orig is None:
raise FileNotFoundError(args.image)
input_np = read_image(args.image, args.height, args.width)
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if 'CUDAExecutionProvider' in ort.get_available_providers() else ['CPUExecutionProvider']
session = ort.InferenceSession(args.onnx, providers=providers)
onnx_output = session.run(['output'], {'input': input_np})[0]
save_comparison(hazy_orig, onnx_output, args.output)
print(f'Saved ONNX result: {args.output}')
if __name__ == '__main__':
main()
|