|
import cv2 |
|
import numpy as np |
|
import axengine as ort |
|
import time |
|
import argparse |
|
|
|
def get_model(model_path: str) -> ort.InferenceSession: |
|
model = ort.InferenceSession(model_path) |
|
|
|
for input in model.get_inputs(): |
|
print(input.name, input.shape, input.dtype) |
|
|
|
for output in model.get_outputs(): |
|
print(output.name, output.shape, output.dtype) |
|
width = model.get_inputs()[0].shape[2] |
|
height = model.get_inputs()[0].shape[1] |
|
|
|
return model, width, height |
|
|
|
|
|
def preprocess_image(image, width=64, height=64): |
|
|
|
h, w = image.shape[:2] |
|
|
|
|
|
scale_ratio = min(width / w, height / h) |
|
|
|
|
|
new_w = int(w * scale_ratio) |
|
new_h = int(h * scale_ratio) |
|
|
|
|
|
resized_img = cv2.resize(image, (new_w, new_h)) |
|
|
|
|
|
letterboxed_img = np.full((height, width, 3), (0, 0, 0), dtype=np.uint8) |
|
|
|
|
|
top = (height - new_h) // 2 |
|
left = (width - new_w) // 2 |
|
|
|
|
|
letterboxed_img[top:top + new_h, left:left + new_w] = resized_img |
|
|
|
|
|
data = np.expand_dims(letterboxed_img, axis=0) |
|
|
|
return data |
|
|
|
|
|
''' |
|
def preprocess_image(image, width=64, height=64): |
|
data = cv2.resize(image, (width, height)) |
|
data = np.expand_dims(data, axis=0) |
|
return data |
|
''' |
|
|
|
if __name__ == "__main__": |
|
parser = argparse.ArgumentParser(description="Process an image with a given model.") |
|
parser.add_argument('--input', type=str, required=True, help='Path to the input image.') |
|
parser.add_argument('--output', type=str, required=True, help='Path to save the output image.') |
|
parser.add_argument('--model', type=str, required=True, help='Path to the model file (.axmodel).') |
|
|
|
args = parser.parse_args() |
|
|
|
model, width, height = get_model(args.model) |
|
|
|
img = cv2.imread(args.input) |
|
|
|
print("Original Image Shape:", img.shape) |
|
img = preprocess_image(img, width, height) |
|
print("Preprocessed Image Shape:", img.shape) |
|
|
|
|
|
start_time = time.time() |
|
|
|
|
|
output = model.run(None, {"input.1": img})[0] |
|
|
|
|
|
end_time = time.time() |
|
elapsed_ms = (end_time - start_time) * 1000 |
|
print(f"Inference Time: {elapsed_ms:.2f} ms") |
|
|
|
print("Output Shape:", output.shape) |
|
|
|
output[output>1] = 1 |
|
output[output<0] = 0 |
|
output_img = (output * 255).astype(np.uint8)[0] |
|
|
|
print("Final Output Image Shape:", output_img.shape) |
|
cv2.imwrite(args.output, output_img) |
|
|