def main(
img_path, dimension, shave, batch_size, scale, print_result, device, model_name
):
"""
Takes an image path and model name and produce the high resolution version
of that image with the help of that model
"""
if model_name == "EDSR":
input_image = ut.load_image(img_path)
elif model_name == "RRDB":
# input_image = ut.load_grayscale_image(img_path)
input_image = ut.npz_loader(img_path)
# print(input_image.shape)
else:
raise Exception("{} : Unknown model...".format(model_name))
output_image = bpfc.patch_batch_forward_chop(
input_image,
dimension,
shave,
scale,
batch_size,
model_type=model_name,
print_timer=True,
)
if print_result:
np.savez("results/outputx4.npz", output_image)
np.save("results/outputx4", output_image)
def main():
# Build a TensorRT engine.
engine = build_engine_onnx("edsr.onnx")
# Inference is the same regardless of which parser is used to build the engine, since the model architecture is the same.
# Allocate buffers and create a CUDA stream.
# input_batch = ut.random_image(100).numpy()
img_path = "data/test2.jpg"
inputs = torch.Tensor(ut.npz_loader(img_path))
# need to set input and output precisions to FP16 to fully enable it
outputs = np.empty([1, 3, 400, 400], dtype=np.float16)
# allocate device memory
d_input = cuda.mem_alloc(1 * inputs.numpy().nbytes)
d_output = cuda.mem_alloc(1 * outputs.nbytes)
bindings = [int(d_input), int(d_output)]
stream = cuda.Stream()
# inputs, outputs, bindings, stream = allocate_buffers(engine)
# Contexts are used to perform inference.
context = engine.create_execution_context()
trt_outputs = do_inference_v2(context,
bindings=bindings,
inputs=inputs,
outputs=outputs,
stream=stream)
print(trt_outputs[0])
print(type(trt_outputs))
print(trt_outputs.shape)
def trt_helper_upsampler_piterative_experiment(model_name,
trt_engine_path,
img_path,
patch_dimension,
shave=10,
use_fp16=False):
"""
Driver function to run a trtengine
Parameters
----------
model_name : str
name of the model (i.e. "EDSR", "RRDB")
trt_engine_path : str
path of the trt engine.
img_path : str
path of the image file.
patch_dimension : int
patch_size.
shave : int, optional
patch overlapping size. The default is 10.
Returns
-------
None.
"""
# Loading model and image
if model_name in ['EDSR']:
img = ut.load_image(img_path)
input_image = img.unsqueeze(0)
elif model_name in ["RRDB"]:
img = ut.npz_loader(img_path)
input_image = img.unsqueeze(0)
else:
print('Unknown model!')
return
b, c, h, w = input_image.shape
total_time = ut.timer()
out_tuple = trt_forward_chop_iterative_v2(
input_image,
trt_engine_path=trt_engine_path,
shave=shave,
min_size=patch_dimension * patch_dimension,
device="cuda",
use_fp16=use_fp16,
print_result=True,
)
output_image = out_tuple[0]
total_time = total_time.toc()
# =============================================================================
# for i in out_tuple[1:]:
# print(i)
# =============================================================================
print('Total executing time: ', total_time)
return output_image
def main(img_path, model_name, patch_dimension):
"""
Driver for recursive forward chop. Takes an image path and model name to upsample the image.
"""
# Loading model and image
img = None
model = None
print("\nLoading model and image... \n")
if model_name in ["EDSR"]:
img = ut.load_image(img_path)
img = img.unsqueeze(0)
model = md.load_edsr(device="cuda")
else:
img = ut.npz_loader(img_path)
img = img.unsqueeze(0)
model = md.load_rrdb(device="cuda")
# Timers for saving stats
timer = [0, 0, 0, 0, 0]
print("Processing...")
# Shfiting input image to CUDA
total_time = ut.timer()
cpu2gpu_time = ut.timer()
img = img.to("cuda")
cpu2gpu_time = cpu2gpu_time.toc()
# Forward chopping and upsampling
output = forward_chop(
img, model, timer=timer, min_size=patch_dimension * patch_dimension
)
# Shifting output image to CPU
gpu2cpu_time = ut.timer()
output = output.to("cpu")
gpu2cpu_time = gpu2cpu_time.toc()
if model_name in ["EDSR"]:
output = output.int()
total_time = total_time.toc()
timer[0] = cpu2gpu_time
timer[-2] = gpu2cpu_time
timer[-1] = total_time
# Saving output
np.savez("results/recursive_outputx4.npz", output)
np.save("results/recursive_outputx4", output)
# Printing processing times
print("\nCPU 2 GPU time: ", timer[0])
print("\nUpsampling time: ", timer[1])
print("\nMerging time: ", timer[2])
print("\nGPU 2 CPU time", timer[3])
print("\nTotal execution time: ", timer[4])
def helper_upsampler_piterative_experiment(model_name, img_path,
patch_dimension):
"""
Driver function for running pytorch model inference
Parameters
----------
img_dimension : int
image one side dimension.
patch_dimension : int
patch size.
Returns
-------
None.
"""
# Loading model and image
if model_name in ['EDSR']:
model = md.load_edsr(device="cuda")
img = ut.load_image(img_path)
input_image = img.unsqueeze(0)
elif model_name in ["RRDB"]:
model = md.load_rrdb(device="cuda")
img = ut.npz_loader(img_path)
input_image = img.unsqueeze(0)
else:
print('Unknown model!')
return
b, c, h, w = input_image.shape
total_time = ut.timer()
out_tuple = forward_chop_iterative(
input_image,
shave=10,
min_size=patch_dimension * patch_dimension,
model=model,
device="cuda",
print_result=True,
)
model.cpu()
del model
output_image = out_tuple[0]
total_time = total_time.toc()
# =============================================================================
# for i in out_tuple[1:]:
# print(i)
# =============================================================================
print('Total executing time: ', total_time)
return output_image
import torch
import modelloader as md
import utilities as ut
model = md.load_rrdb(device="cuda")
img = ut.npz_loader("data/slices/0.npz")
img = img[:, 0:100, 0:100]
img = img.unsqueeze(0)
img = img.to("cuda")
model.eval()
with torch.no_grad():
output = model(img)
def main(stat_path, model_name, img_path, shave, scale):
total_patches = "Total Patches"
total_batches_ = "Total Batches"
maximum_batch_size = "Maximum Batch Size"
total_batch_processing_time = "Total batch processing time"
per_batch_processing_time_ = "Per Batch Processing Time"
patch_dimension = "Patch Dimension"
total_time = "Total time"
stat_df = pd.read_csv(stat_path)
# =============================================================================
# print(stat_df.columns)
# =============================================================================
total_batches = stat_df[total_patches] / stat_df[maximum_batch_size]
stat_df[total_batches_] = total_batches
per_batch_processing_time = (
stat_df[total_batch_processing_time] / stat_df[total_batches_]
)
stat_df[per_batch_processing_time_] = per_batch_processing_time
# =============================================================================
# print(stat_df.columns)
# print(stat_df)
# =============================================================================
maximum_patch_size = stat_df[patch_dimension].max()
min_total_processing_time = stat_df[total_time].min()
idx_min_total_processing_time = stat_df[total_time].idxmin()
min_per_batch_processing_time = stat_df[per_batch_processing_time_].min()
idx_min_per_batch_processing_time = stat_df[per_batch_processing_time_].idxmin()
print(min_total_processing_time)
print(idx_min_total_processing_time)
print(min_per_batch_processing_time)
print(idx_min_per_batch_processing_time)
print("here")
# =============================================================================
# patch_dimension = 0
# batch_size = 0
# =============================================================================
# =============================================================================
# print(stat_df.loc[idx_min_total_processing_time, patch_dimension])
# =============================================================================
img = ut.npz_loader(img_path)
c, h, w = img.shape
if h < w:
temp = h
h = w
w = temp
dimensions = stat_df.loc[
:, [patch_dimension, maximum_batch_size, per_batch_processing_time_]
].values
total_patches_from_img = []
for d in range(len(dimensions)):
img_patch_count = ut.patch_count(h, w, dimensions[d][0], shave)
img_batch_count = img_patch_count / dimensions[d][1]
img_processing_time = img_batch_count * dimensions[d][2]
total_patches_from_img.append(
[
dimensions[d][0],
dimensions[d][1],
img_patch_count,
math.ceil(img_batch_count),
img_processing_time,
]
)
# print(total_patches_from_img[269])
img_df = pd.DataFrame(total_patches_from_img)
best_index = img_df[4].idxmin()
# =============================================================================
# print(best_index)
# print(img_df.iloc[best_index, 0])
# =============================================================================
if h < maximum_patch_size and w < maximum_patch_size:
patch_dimension = int(img_df.iloc[idx_min_total_processing_time, 0])
batch_size = int(img_df.iloc[idx_min_total_processing_time, 1])
else:
patch_dimension = int(img_df.iloc[best_index, 0])
batch_size = int(img_df.iloc[best_index, 1])
print("Patch dimension: {}, batch size: {}".format(patch_dimension, batch_size))
return engine
engine = build_engine("edsr.onnx")
with open("edsr.trt", "wb") as f:
f.write(bytearray(engine.serialize()))
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(TRT_LOGGER)
with open("edsr.trt", "rb") as f:
engine_bytes = f.read()
engine = runtime.deserialize_cuda_engine(engine_bytes)
edsr_context = engine.create_execution_context()
# input_batch = ut.random_image(100).numpy()
img_path = "data/test2.jpg"
input_batch = ut.npz_loader(img_path).numpy()
# need to set input and output precisions to FP16 to fully enable it
output_d = np.empty([1, 3, 400, 400], dtype=np.float32)
"""
memory allocation for inputs
"""
d_input = cuda.mem_alloc(1 * input_batch.nbytes)
"""memory allocation for outputs"""
d_output = cuda.mem_alloc(1 * output_d.nbytes)
bindings = [int(d_input), int(d_output)]
stream = cuda.Stream()
def predict(batch): # result gets copied into output
# transfer input data to device
# Saving loader config for resuming the process if needed
loader_config_file = open("../loader_config.toml", "w")
toml.dump(loader_config, loader_config_file)
loader_config_file.close()
# Channels, height, and width of the image
c, h, w = 0, 0, 0
if model_name not in ["RRDB"]:
# Loads normal JPEG image
img = ut.load_image(img_path)
c, h, w = img.shape
else:
# Loads special image for models like RRDB...
img = ut.npz_loader(img_path)
c, h, w = img.shape
# If the maximum dimension of the patch is greater than images height
# or width then stop the process
if max_dim > h or max_dim > w:
raise Exception(
"end_patch_dimension in batch_processing.toml is greater than \
input image dimension. Use a bigger input image or change end_patch_dimension. "
)
# Batch range checker process
full_result = batch_range_checker(
max_dim,
min_dim,
shave,
def upsample(model_name, img_path, dimension, shave, batch_size, scale,
device):
file_name = img_path.split("/")[-1].split(".")[0]
if model_name == "RRDB":
input_image = ut.npz_loader(img_path)
c, h, w = input_image.shape
patch_list = {}
create_patch_list(patch_list, input_image, dimension, shave, scale, c,
h, w)
model = md.load_rrdb(device=device)
model.eval()
min_dim = min(dimension, h, w)
if min_dim != dimension:
print(
"\nPatch dimension is greater than the input image's minimum dimension. Changing patch dimension to input image's minimum dimension... \n "
)
dimension = min_dim
output, _ = batch_forward_chop(
patch_list,
batch_size,
c,
h,
w,
dimension,
shave,
scale,
model=model,
device=device,
)
output = output.int()
output_folder = "output_images"
ut.save_image(output,
output_folder,
h,
w,
scale,
output_file_name=file_name + "_x4")
elif model_name == "EDSR":
input_image = ut.load_image(img_path)
c, h, w = input_image.shape
patch_list = {}
create_patch_list(patch_list, input_image, dimension, shave, scale, c,
h, w)
model = md.load_edsr(device=device)
model.eval()
min_dim = min(dimension, h, w)
if min_dim != dimension:
print(
"\nPatch dimension is greater than the input image's minimum dimension. Changing patch dimension to input image's minimum dimension... \n "
)
dimension = min_dim
output, _ = batch_forward_chop(
patch_list,
batch_size,
c,
h,
w,
dimension,
shave,
scale,
model=model,
device=device,
)
# =============================================================================
# print(output)
# =============================================================================
output = output.int()
output_folder = "output_images"
ut.save_image(output,
output_folder,
h,
w,
scale,
output_file_name=file_name + "_x4")
else:
print("Unknown model...")
onnx_model = onnx.load("inference_models/rrdb.onnx")
onnx.checker.check_model(onnx_model)
import onnxruntime
ort_session = onnxruntime.InferenceSession("inference_models/rrdb.onnx")
def to_numpy(tensor):
return (tensor.detach().cpu().numpy()
if tensor.requires_grad else tensor.cpu().numpy())
img_path = "data/slices/13.npz"
input_batch = ut.npz_loader(img_path).unsqueeze(0)
input_ = torch.tensor(input_batch).int()
output_folder = "output_images"
file_name = img_path.split("/")[-1].split(".")[0]
ut.save_image(input_[0],
output_folder,
30,
30,
4,
output_file_name=file_name + "_input_x4")
print(input_batch.shape)
ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(input_batch)}
execution_time = ut.timer()