def binary_search_helper(dimension, logger, model_name="EDSR", device="cuda"):
"""
Process random image and calculates processing time
Parameters
----------
dimension : int
random image dimension.
logger : logger
keep logs.
device : str, optional
GPU or CPU. The default is 'cuda'.
Returns
-------
total_time : float
EDSR processing time.
"""
print('Before loading model: ')
subprocess.run("gpustat", shell=True)
print()
total_time = 0
try:
model = None
if model_name == "EDSR":
model = md.load_edsr(device=device)
print('After loading model: ')
subprocess.run("gpustat", shell=True)
print()
elif model_name == "RRDB":
model = md.load_rrdb(device=device)
else:
raise Exception("Unknown model...")
model.eval()
input_image = ut.random_image(dimension)
if model_name == "RRDB":
input_image = input_image[:, 2:, :, :]
input_image = input_image.to(device)
with torch.no_grad():
start = time.time()
print('Before processing: ')
subprocess.run("gpustat", shell=True)
output_image = model(input_image)
print('After processing: ')
subprocess.run("gpustat", shell=True)
end = time.time()
total_time = end - start
ut.clear_cuda(input_image, output_image)
model.cpu()
del model
print('After model shifting and deleting: ')
subprocess.run("gpustat", shell=True)
except RuntimeError as err:
logger.error("Runtime error for dimension: {}x{}: " + err)
sys.exit(1)
return total_time
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_rrdb_experiment(img_dimension, patch_dimension):
# Loading model and image
img = None
model = None
# =============================================================================
# print("\nLoading model and image... \n")
# =============================================================================
img = np.load("data/slices/0.npz")
img = img.f.arr_0
img = np.resize(img, (img_dimension, img_dimension))
img = img[np.newaxis, :, :]
img = torch.from_numpy(img)
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()
total_time = total_time.toc()
timer[0] = cpu2gpu_time
timer[-2] = gpu2cpu_time
timer[-1] = total_time
# 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])
# =============================================================================
print(timer[1])
print(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
def helper_rrdb_piterative_experiment(img_dimension, 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
img = None
model = None
img = np.load("data/slices/0.npz")
img = img.f.arr_0
img = np.resize(img, (img_dimension, img_dimension))
img = img[np.newaxis, :, :]
img = torch.from_numpy(img)
img = img.unsqueeze(0)
input_image = img
b, c, h, w = input_image.shape
# input_image = input_image.reshape((1, c, h, w))
# Loading model
model = md.load_rrdb("cuda")
model.eval()
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_time)
def main(model_name, onnx_model_name):
device = ut.get_device_type()
model = None
if model_name == "RRDB":
model = md.load_rrdb(device)
dummy_input = ut.random_image(100, batch=False, channel=1).to(device)
elif model_name == "EDSR":
model = md.load_edsr(device)
dummy_input = ut.random_image(100).to(device)
model.eval()
torch.onnx.export(model,
dummy_input,
"inference_models/" + onnx_model_name,
verbose=False)
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 helper_rrdb_iterative(stat_path, img_dimension, 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 = np.load("data/slices/0.npz")
img = img.f.arr_0
img = np.resize(img, (img_dimension, img_dimension))
img = img[np.newaxis, :, :]
img = torch.from_numpy(img)
# =============================================================================
# img = img.unsqueeze(0)
# =============================================================================
model = md.load_rrdb(device="cuda")
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))
# =============================================================================
# =============================================================================
# bpfc.upsample(
# model_name,
# img_path,
# patch_dimension,
# shave,
# batch_size,
# scale,
# device="cuda",
# )
# =============================================================================
# =============================================================================
# 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))
# =============================================================================
print(patch_dimension)
print(batch_size)
bpfc.patch_batch_forward_chop(
input_image=img,
patch_dimension=patch_dimension,
patch_shave=shave,
scale=scale,
batch_size=batch_size,
model_type="RRDB",
device="cuda",
print_timer=True,
)
def build_onnx_model(model_name, patch_size, onnx_model_name, device="cuda"):
"""
Builds ONNX model with a fixed input shape
Parameters
----------
model_name : str
Upsampler model name. (i.e. RRDB, EDSR).
patch_size : int
Input image dimension n of nxn.
onnx_model_name : str
output ONNX model name.
Returns
-------
None.
"""
# =============================================================================
# print('Before loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# total_time = 0
# model = None
# if model_name == "EDSR":
# model = md.load_edsr(device=device)
# print('After loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# elif model_name == "RRDB":
# model = md.load_rrdb(device=device)
# else:
# raise Exception("Unknown model...")
# model.eval()
# input_image = ut.random_image(patch_size)
# if model_name == "RRDB":
# input_image = input_image[:, 2:, :, :]
# dummy_input = input_image.to(device)
# =============================================================================
# =============================================================================
# print('Before loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
device = ut.get_device_type()
model = None
if model_name == "RRDB":
model = md.load_rrdb(device)
image = ut.create_custom_npz(patch_size, patch_size)
image = image[np.newaxis, :, :]
image = torch.from_numpy(image)
dummy_input = image.unsqueeze(0).to(device)
elif model_name == "EDSR":
model = md.load_edsr(device)
# =============================================================================
# print('After loading model: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
dummy_input = ut.random_image(patch_size).to(device)
# =============================================================================
# print('After loading image: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
# =============================================================================
# b, c, h, w = 1, 3, patch_size, patch_size
# dummy_input = torch.rand(b, c, h, w, requires_grad=False).to(device)
# =============================================================================
else:
print("Unknown model!")
return
print(dummy_input.shape)
model.eval()
with torch.no_grad():
# =============================================================================
# print('Before processing: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
# =============================================================================
# output = model(input_image )
# print('After processing: ')
# subprocess.run("gpustat", shell = True)
# print()
# =============================================================================
try:
torch.onnx.export(
model,
dummy_input,
"inference_models/" + onnx_model_name,
verbose=False,
opset_version=12,
input_names=["input"],
output_names=["output"],
)
except RuntimeError as err:
sys.exit(1)
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...")
def patch_batch_forward_chop(
input_image,
patch_dimension,
patch_shave,
scale,
batch_size,
model_type="EDSR",
device="cuda",
print_timer=True,
):
"""
Parameters
----------
input_image : 3D Matrix
input image.
patch_dimension : int
patch dimension.
patch_shave : int
patch shave value.
scale : int
scale for LR to SR.
batch_size : int
batch size.
model_type : str, optional
model name. The default is 'EDSR'.
device : str, optional
GPU or CPU. The default is 'cuda'.
print_timer : bool, optional
print result or not. The default is True.
Returns
-------
output_image : 3D Matrix
output SR image.
"""
model = None
if model_type == "EDSR":
model = md.load_edsr(device=device)
model.eval()
elif model_type == "RRDB":
model = md.load_rrdb(device=device)
model.eval()
else:
raise Exception("{} : Unknown model...".format(model_type))
total_timer = ut.timer()
channel, height, width = input_image.shape
patch_list_timer = ut.timer()
patch_list = {}
create_patch_list(
patch_list,
input_image,
patch_dimension,
patch_shave,
scale,
channel,
height,
width,
)
patch_list_processing_time = patch_list_timer.toc()
total_batch_processing_timer = ut.timer()
output_image, timer_results = batch_forward_chop(
patch_list,
batch_size,
channel,
height,
width,
patch_dimension,
patch_shave,
scale,
model=model,
device="cuda",
print_timer=False,
)
total_batch_processing_time = total_batch_processing_timer.toc()
if model_type == "EDSR":
output_image = output_image.int()
total_time = total_timer.toc()
print(len(patch_list))
if print_timer:
print(patch_list_processing_time)
for t in timer_results:
print(t)
print(total_batch_processing_time)
print(total_time)
model = model.cpu()
del model
return output_image
img = ut.npz_loader(img_path)
c, h, w = img.shape
# img = img.reshape((1, c, h, w))
# plt.imshow(img[0].permute((1,2,0)))
input_image = img.float()
# Creating patch list from the image
patch_list_timer = ut.timer()
patch_list = {}
create_patch_list(patch_list, input_image, dimension, shave, 4, c, h, w)
patch_list_processing_time = patch_list_timer.toc()
print("Total patch list creating time: {}".format(
patch_list_processing_time))
print(len(patch_list))
# Loading model
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
# Batch Processing
batch_processing_start = time.time()
output, _ = batch_forward_chop(patch_list,
batch_size,
c,
h,