File size: 2,260 Bytes
93e8fc6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
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)
|