import torch
import cv2
import argparse

import numpy as np

if __name__=='__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--model', '-m', type = str, required = True, help = "Path to Torchscript \
      model.")
    parser.add_argument('--input_image', '-i', type = str, required = True, help = "Path to input \
      image.")
    parser.add_argument('--output_image', '-o', type = str, required = True, help = "Path to output \
      image.")
    parser.add_argument('--use_gpu', '-g', type = bool, default = True)

    args = parser.parse_args()

    model = torch.jit.load(args.model)
    if args.use_gpu:
        model = model.half()
        model = model.cuda()

    img = cv2.imread(args.input_image)
    img = img[:, :, [2, 1, 0]] / 255.0
    img = torch.from_numpy(img)
    img = img.permute([2, 0, 1]).unsqueeze(0)
    if args.use_gpu:
        img = img.half().cuda()
    with torch.no_grad():
      output = model(img)

    output = output.cpu().float().squeeze(0)
    output = output.permute([1, 2, 0])
    output = output.numpy()
    output = output[: , :, [2, 1, 0]]
    output = np.clip(output, 0, 1.0) * 255.0

    cv2.imwrite(args.output_image, output.astype(np.uint8))
    print("Done!!")
