AXERA-TECH/CodeFormer

1import argparse2import os3import cv24import numpy as np5import axengine as axe6 7def from_numpy(x):8    return x if isinstance(x, np.ndarray) else np.array(x)9 10def main(args):11    # Initialize the model12    session = axe.InferenceSession(args.model_path)13    output_names = [x.name for x in session.get_outputs()]14    input_name = session.get_inputs()[0].name15 16    # results17    os.makedirs(args.output_path, exist_ok=True)18 19    files =[f for f in os.listdir(args.inputs_path) if f.lower().endswith(('.jpg', '.png', 'jpeg'))]20    21    for file in files:22        ori_image = cv2.imread(os.path.join(args.inputs_path, file))23        h, w = ori_image.shape[:2]24        image = cv2.resize(ori_image, (512, 512))25        image = (image[..., ::-1] /255.0).astype(np.float32)26        27        mean = [0.5, 0.5, 0.5]28        std = [0.5, 0.5, 0.5]29        image = ((image - mean) / std).astype(np.float32)30 31        #image = (image /1.0).astype(np.float32)32        img = np.transpose(np.expand_dims(np.ascontiguousarray(image), axis=0), (0,3,1,2))33        34        # Use the model to generate super-resolved images35        sr = session.run(output_names, {input_name: img})36 37        #sr_y_image = imgproc.array_to_image(sr)38        sr = np.transpose(sr[0].squeeze(0), (1,2,0))39        sr = (sr*std + mean).astype(np.float32)40        41        # Save image42        ndarr = np.clip((sr*255.0), 0, 255.0).astype(np.uint8)43        out_image = cv2.resize(ndarr[..., ::-1], (w, h))44 45        cv2.imwrite(f'{args.output_path}/{file}', out_image)46        print(f"SR image save to `{file}`")47 48 49if __name__ == "__main__":50    parser = argparse.ArgumentParser(description="Using the model generator super-resolution images.")51    parser.add_argument("--inputs_path",52                        type=str,53                        default="images",54                        help="origin image path.")55    parser.add_argument("--output_path",56                        type=str,57                        default="results",58                        help="colorized image path.")59    parser.add_argument("--model_path",60                        type=str,61                        default="./codeformer.axmoel",62                        help="model path.")63    args = parser.parse_args()64 65    main(args)66