File size: 2,918 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 | import argparse
import os
import numpy as np
import onnxruntime as ort
from PIL import Image
HEIGHT = 512
WIDTH = 512
IMG_EXTENSIONS = ('.jpg', '.JPG', '.jpeg', '.JPEG', '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP')
def parse_args():
parser = argparse.ArgumentParser(description='Run GCANet ONNX inference with fixed 512x512 input.')
parser.add_argument('--task', default='dehaze', choices=['dehaze', 'derain'])
parser.add_argument('--onnx', default=None, help='Path to ONNX model. Default: onnx/gcanet_{task}_512x512_sim.onnx')
parser.add_argument('--indir', default='examples')
parser.add_argument('--outdir', default='onnx_output')
return parser.parse_args()
def make_dataset(image_dir):
images = []
assert os.path.isdir(image_dir), '%s is not a valid directory' % image_dir
for root, _, fnames in sorted(os.walk(image_dir)):
for fname in fnames:
if fname.endswith(IMG_EXTENSIONS):
images.append(os.path.join(root, fname))
return images
def edge_compute_np(img_chw):
x_diffx = np.abs(img_chw[:, :, 1:] - img_chw[:, :, :-1])
x_diffy = np.abs(img_chw[:, 1:, :] - img_chw[:, :-1, :])
edge = np.zeros_like(img_chw, dtype=np.float32)
edge[:, :, 1:] += x_diffx
edge[:, :, :-1] += x_diffx
edge[:, 1:, :] += x_diffy
edge[:, :-1, :] += x_diffy
edge = np.sum(edge, axis=0, keepdims=True) / 3.0
edge = edge / 4.0
return edge.astype(np.float32)
def preprocess(img_path):
img = Image.open(img_path).convert('RGB')
img = img.resize((WIDTH, HEIGHT), Image.BICUBIC)
img_np = np.array(img).astype(np.float32)
img_chw = np.transpose(img_np, (2, 0, 1))
edge = edge_compute_np(img_chw)
model_input = np.concatenate((img_chw, edge), axis=0)[None, :, :, :] - 128.0
return img_chw, model_input.astype(np.float32)
def postprocess(pred, img_chw, only_residual):
out = pred[0]
if only_residual:
out = out + img_chw
out = np.round(out).clip(0, 255).astype(np.uint8)
out = np.transpose(out, (1, 2, 0))
return out
def main():
args = parse_args()
onnx_path = args.onnx or os.path.join('onnx', 'gcanet_%s_512x512_sim.onnx' % args.task)
only_residual = args.task == 'dehaze'
os.makedirs(args.outdir, exist_ok=True)
session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
input_name = session.get_inputs()[0].name
for img_path in make_dataset(args.indir):
img_chw, model_input = preprocess(img_path)
pred = session.run(None, {input_name: model_input})[0]
out_img = postprocess(pred, img_chw, only_residual)
save_name = os.path.splitext(os.path.basename(img_path))[0] + '_%s_onnx.png' % args.task
Image.fromarray(out_img).save(os.path.join(args.outdir, save_name))
print('Saved:', os.path.join(args.outdir, save_name))
if __name__ == '__main__':
main()
|