import argparse
import os
import cv2
import numpy as np
import axengine as axe

def from_numpy(x):
    return x if isinstance(x, np.ndarray) else np.array(x)

def main(args):
    # Initialize the model
    session = axe.InferenceSession(args.model_path)
    output_names = [x.name for x in session.get_outputs()]
    input_name = session.get_inputs()[0].name

    # results
    os.makedirs(args.output_path, exist_ok=True)

    files =[f for f in os.listdir(args.inputs_path) if f.lower().endswith(('.jpg', '.png', 'jpeg'))]
    
    for file in files:
        ori_image = cv2.imread(os.path.join(args.inputs_path, file))
        h, w = ori_image.shape[:2]
        image = cv2.resize(ori_image, (512, 512))
        image = (image[..., ::-1] /255.0).astype(np.float32)
        
        mean = [0.5, 0.5, 0.5]
        std = [0.5, 0.5, 0.5]
        image = ((image - mean) / std).astype(np.float32)

        #image = (image /1.0).astype(np.float32)
        img = np.transpose(np.expand_dims(np.ascontiguousarray(image), axis=0), (0,3,1,2))
        
        # Use the model to generate super-resolved images
        sr = session.run(output_names, {input_name: img})

        #sr_y_image = imgproc.array_to_image(sr)
        sr = np.transpose(sr[0].squeeze(0), (1,2,0))
        sr = (sr*std + mean).astype(np.float32)
        
        # Save image
        ndarr = np.clip((sr*255.0), 0, 255.0).astype(np.uint8)
        out_image = cv2.resize(ndarr[..., ::-1], (w, h))

        cv2.imwrite(f'{arg.output_path}/{file}', out_image)
        print(f"SR image save to `{file}`")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Using the model generator super-resolution images.")
    parser.add_argument("--inputs_path",
                        type=str,
                        default="images",
                        help="origin image path.")
    parser.add_argument("--output_path",
                        type=str,
                        default="results",
                        help="colorized image path.")
    parser.add_argument("--model_path",
                        type=str,
                        default="./codeformer.axmoel",
                        help="model path.")
    args = parser.parse_args()

    main(args)