import time

import cv2
import numpy as np
import torch
import tritonclient.http as httpclient
from PIL import Image

triton_client = httpclient.InferenceServerClient(url=f"10.1.1.150:7000")

sample = cv2.imread("comic2.png", cv2.IMREAD_COLOR).astype(np.float32) / 255.
sample = np.transpose(sample if sample.shape[2] == 1 else sample[:, :, [2, 1, 0]], (2, 0, 1))
sample = torch.from_numpy(sample).float().unsqueeze(0).numpy()
inputs = [
    httpclient.InferInput("input", sample.shape, datatype="FP32")
]
inputs[0].set_data_from_numpy(sample, binary_data=True)
start_time = time.time()
results = triton_client.infer(model_name="super_resolution", inputs=inputs)
print(time.time() - start_time)
sr_output = torch.from_numpy(results.as_numpy(f"output"))
output = sr_output.data.squeeze().float().cpu().clamp_(0, 1).numpy()
if output.ndim == 3:
    output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))  # CHW-RGB to HCW-BGR
output = (output * 255.0).round().astype(np.uint8)
# cv2.imshow("", output)
# cv2.waitKey(0)
cv2.imwrite("comic3.png", output)