import numpy as np
import cv2
import os
from hobot_dnn import pyeasy_dnn as dnn
from realesrgan import RealESRGANer
import argparse
from PIL import Image
import torch

def get_hw(pro):
    if pro.layout == "NCHW":
        return pro.shape[2], pro.shape[3]
    else:
        return pro.shape[1], pro.shape[2]

modelpath_prefix = ''

# img_path 图像完整路径
img_path = '/root/Downloads/Real-ESRGAN-master/inputs/ADE_val_00000114.jpg'
# model_path 量化模型完整路径
model_path = '/root/Downloads/Real-ESRGAN-master/weights/RealESRGAN_x4plus_anime_6B.bin'

# 1. 加载模型，获取所需输出HW
models = dnn.load(model_path)
model_h, model_w = get_hw(models[0].inputs[0].properties)
print(model_h, model_w)
# 2 加载图像，根据前面模型，转换后的模型是以NV12作为输入的
# 但在OE验证的时候，需要将图像再由NV12转为YUV444
imgOri = cv2.imread(img_path,cv2.IMREAD_UNCHANGED)#248,354,3
h_input, w_input = imgOri.shape[0:2]
print("imgOri.dtype",imgOri.dtype)
# imgOri = imgOri.astype(np.float32) / 255
img = cv2.resize(imgOri, (model_w, model_h))#HWC,224,224,3


# img=torch.from_numpy(np.transpose(img, (2, 0, 1)))#224,224,3
img=torch.from_numpy(img)
img = img.unsqueeze(0)#1,224,224,3

# 3 模型推理
def print_properties(pro):
    print("tensor type:", pro.tensor_type)
    print("data type:", pro.dtype)
    print("layout:", pro.layout)
    print("shape:", pro.shape)
print_properties(models[0].inputs[0].properties)
for output in models[0].outputs:
    print_properties(output.properties)


outputs = models[0].forward(img)
# print("outputs[0].buffer的shape是",np.shape(outputs[0].buffer))#NCHW(1, 3, 896, 896)
# print("outputs[0].buffer[0]的shape是",np.shape(outputs[0].buffer[0]))#CHW(3, 896, 896)
# print("outputs[0].buffer[0][0]的shape是",np.shape(outputs[0].buffer[0][0]))#(896, 896)
# print("outputs[0].buffer[0]的类型是",type(outputs[0].buffer[0]))#numpy.ndarray
# print("outputs[0].buffer[0]的shape是",np.shape(outputs[0].buffer[0]))#(3,896, 896)
# print("outputs[0].buffer[0]是",outputs[0].buffer[0])
rimg0 = np.transpose(3000*(outputs[0].buffer[0]), [1, 2, 0]).round().astype(np.uint8)
print("rimg0是",rimg0)
cv2.imwrite('/root/Downloads/Real-ESRGAN-master/0.jpg',rimg0)#绿色
print("保存成功")
