<jupyter_start><jupyter_text>Checking and testing inference<jupyter_code># from models.psp_project.helpers.common import detect_and_align_face
# from onnx_tf.backend import prepare

# https://stackoverflow.com/questions/67480507/tensorflow-equivalent-of-pytorchs-transforms-normalize
def normalize_image_pytorch_style(image, mean, std):
    for channel in range(3):
        image[:,:,channel] = (image[:,:,channel] - mean[channel]) / std[channel]
    return image



def load_preprocess_image_from_path(img_path):
    print("Image", img_path)
    
    # Open the image into the right dimensions
    img = Image.open(img_path)
    img = img.convert("RGB")
    img = img.resize((256, 256))

    # Normalized img dimensions to between -1 and 1
    # similar to how the original model did this in pytorch
    img_resized = normalize_image_pytorch_style(
                        np.array(img) / 255.0, 
                        mean=[0.5, 0.5, 0.5], 
                        std=[0.5, 0.5, 0.5])

    
    # Move color axis to start like model is used to
    img_resized = np.moveaxis(img_resized, 2, 0)

    # Add in the first batch dim
    input_batch_img = np.expand_dims(img_resized, axis=0)
    
    return input_batch_img.astype(np.float32)
tic = time.time()

# Convert model to tensorflow
model_psp = onnx.load(onnx_path)
# tf_rep_psp = prepare(model_psp)

toc = time.time()
print('Loading and converting PSP onnx to tf took {:.4f} seconds.'.format(toc - tic))
test_img = "data/portraits/dongyun.jpg"

ort_session = ort.InferenceSession(onnx_path)
input_name = ort_session.get_inputs()[0].name
ouput_name = ort_session.get_outputs()[0].name
print(f"ONNX Input name -> {input_name} and output -> {ouput_name}")

img = load_preprocess_image_from_path(test_img)



output = ort_session.run([ouput_name],
                        {input_name: img})[0]
output

# 
# print(output)

output = (output.squeeze().transpose((1, 2, 0)) + 1) / 2
output[output < 0] = 0
output[output > 1] = 1
output = normalization(output) * 255
output = Image.fromarray(output.astype('uint8'))
output

# output.save(os.path.join(args.save, "{}_{}_{}.png".format(
#     args.network,
#     args.platform,
#     image.split("/")[-1].split(".")[0]))
# )
out_img, out_lat = output

print(output)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13219895287958114, 'fraction_numerical': 0.024869109947643978, 'mean_word_length': 3.5496031746031744, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 8, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '486473', 'license': 'MIT', 'n_tokens_mistral': 851, 'n_tokens_neox': 801, 'n_words': 171, 'path': 'v3_psp/convert_torch_to_onnx.ipynb', 'repo_name': 'vazkir/realworld-stylegan2-encoder', 'size': '2292'}