AXERA-TECH/Real-ESRGAN

1import cv22import numpy as np3import axengine as ort4import time5import argparse6 7def get_model(model_path: str) -> ort.InferenceSession:8    model = ort.InferenceSession(model_path)9 10    for input in model.get_inputs():11        print(input.name, input.shape, input.dtype)12 13    for output in model.get_outputs():14        print(output.name, output.shape, output.dtype)15    width = model.get_inputs()[0].shape[2]16    height = model.get_inputs()[0].shape[1]17    18    return model, width, height19 20 21def preprocess_image(image, width=64, height=64):22    # 获取原始图像的高度和宽度23    h, w = image.shape[:2]24    25    # 计算调整大小的比例26    scale_ratio = min(width / w, height / h)27    28    # 根据比例计算新的高度和宽度，同时保持原图宽高比29    new_w = int(w * scale_ratio)30    new_h = int(h * scale_ratio)31    32    # 调整图像大小，保持原图宽高比33    resized_img = cv2.resize(image, (new_w, new_h))34    35    # 创建一个具有目标尺寸的空白图像（黑色背景）36    letterboxed_img = np.full((height, width, 3), (0, 0, 0), dtype=np.uint8)37    38    # 计算将调整大小后的图像放置在letterbox中的起始点39    top = (height - new_h) // 240    left = (width - new_w) // 241    42    # 将调整大小后的图像放入letterboxed图像中43    letterboxed_img[top:top + new_h, left:left + new_w] = resized_img44    45    # 添加批次维度46    data = np.expand_dims(letterboxed_img, axis=0)47    48    return data49 50 51'''52def preprocess_image(image, width=64, height=64):53    data = cv2.resize(image, (width, height))54    data = np.expand_dims(data, axis=0)55    return data56'''57 58if __name__ == "__main__":59    parser = argparse.ArgumentParser(description="Process an image with a given model.")60    parser.add_argument('--input', type=str, required=True, help='Path to the input image.')61    parser.add_argument('--output', type=str, required=True, help='Path to save the output image.')62    parser.add_argument('--model', type=str, required=True, help='Path to the model file (.axmodel).')63 64    args = parser.parse_args()65 66    model, width, height = get_model(args.model)67    68    img = cv2.imread(args.input)69    70    print("Original Image Shape:", img.shape)71    img = preprocess_image(img, width, height)72    print("Preprocessed Image Shape:", img.shape)73    74    # 开始计时75    start_time = time.time()76 77    # 执行推理78    output = model.run(None, {"input.1": img})[0]79 80    # 结束计时并计算耗时（毫秒）81    end_time = time.time()82    elapsed_ms = (end_time - start_time) * 1000  # 秒转毫秒83    print(f"Inference Time: {elapsed_ms:.2f} ms")84 85    print("Output Shape:", output.shape)86    87    output[output>1] = 188    output[output<0] = 089    output_img = (output * 255).astype(np.uint8)[0]90    91    print("Final Output Image Shape:", output_img.shape)92    cv2.imwrite(args.output, output_img)93