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