When call onnx model, how to convert image file to correct model input

Question

I try this repo: https://github.com/cszn/KAIR

I use following code build onnx

import os.path

import torch

# Define your model architecture and load its weights
from models.network_srmd import SRMD as net

model_name = 'srmdnf_x4'  # Specify the model name
model = net(in_nc=18, out_nc=3, nc=128, nb=12, upscale=4, act_mode='R', upsample_mode='pixelshuffle')

# Load the model weights
pj = os.path.abspath(f"{os.path.dirname(__file__)}/..")
model_path = f"{pj}/model_zoo/{model_name}.pth"
model.load_state_dict(torch.load(model_path))

# Set the model to evaluation mode
model.eval()

# Provide sample input dimensions (change according to your model's input requirements)
input_shape = (1, 18, 256, 256)

# Export the model to ONNX format
dummy_input = torch.randn(input_shape)
output_path = f"{pj}/model_onnx/{model_name}.onnx"
torch.onnx.export(model, dummy_input, output_path, verbose=True)

print('ONNX model exported successfully to:', output_path)

when I use following code call onnx model

import os.path
import onnxruntime as ort
import numpy as np
from PIL import Image

# Load the ONNX model
d = os.path.dirname(__file__)
model_path = f'{d}/srmdnf_x4.onnx'
ort_session = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])

# Get input and output names
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name

# Load the input image
input_image_path = f'{d}/t.jpg'
input_image = Image.open(input_image_path).convert('RGB')
input_image = input_image.resize((256, 256))  # Resize image to match model input size

# Preprocess the input image
input_image = np.array(input_image).astype(np.float32) / 255.0
input_image = np.transpose(input_image, (2, 0, 1))  # Change the image shape to (C, H, W)
input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension

# Run inference with the ONNX model
outputs = ort_session.run([output_name], {input_name: input_image})

# Postprocess the output to get the final image
output_image = outputs[0][0]
output_image = np.clip(output_image * 255.0, 0, 255).astype(np.uint8)
output_image = Image.fromarray(np.transpose(output_image, (1, 2, 0)))  # Change the image shape back to (H, W, C)

# Save the output image
output_image.save(f'{d}/out.jpg')
print('Output image saved as out.jpg')

it show err

test_srmd.py:None (test_srmd.py)
test_srmd.py:26: in <module>
    outputs = ort_session.run([output_name], {input_name: input_image})
../../../../.pyenv/versions/3.9.1/lib/python3.9/site-packages/onnxruntime/capi/onnxruntime_inference_collection.py:220: in run
    return self._sess.run(output_names, input_feed, run_options)
E   onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument: [ONNXRuntimeError] : 2 : INVALID_ARGUMENT : Got invalid dimensions for input: input.1 for the following indices
E    index: 1 Got: 3 Expected: 18
E    Please fix either the inputs or the model.
collected 0 items / 1 error

I just confused, changing build onnx code or chaning call onnx code, which one should i do to make onnx work well