From cc6e5490795174965e982002a77ea25917ba2612 Mon Sep 17 00:00:00 2001 From: Nate Date: Thu, 21 May 2020 14:39:49 -0500 Subject: [PATCH] Removed prints --- ColabESRGAN/test.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/ColabESRGAN/test.py b/ColabESRGAN/test.py index e54fbd3..808c472 100644 --- a/ColabESRGAN/test.py +++ b/ColabESRGAN/test.py @@ -4,6 +4,7 @@ import numpy as np import torch import architecture +import math # ESRGAN class allows abstraction of warmup and inference. class esrgan(): @@ -24,25 +25,25 @@ def __init__(self, model_path=None, hw='cpu'): print('Model warmup complete') # Function to run esrgan on single image, and single output. - def run_esrgan(self, test_img_folder=None, out_filename=None): + def run_esrgan(self, test_img_folder=None, out_filename=None, mosaic_res=1): assert out_filename assert test_img_folder - # test_img_folder = 'LR/*' - # print('Model path {:s}. \nTesting...'.format(model_path)) - - # idx = 0 - # for path in glob.glob(test_img_folder): - # idx += 1 - # base = os.path.splitext(os.path.basename(path))[0] - # print(idx, base) - # read image - # print("Running ESRGAN on ", test_img_folder) - # print(test_img_folder) img = cv2.imdecode(np.fromfile(test_img_folder, np.uint8), cv2.IMREAD_UNCHANGED) + # from alex: codes to shrink image if memory is an issue + # GPUmem = torch.cuda.get_device_properties(0).total_memory + # Sx = int(1.2*img.shape[0]/mosaic_res) + # Sy = int(1.2*img.shape[1]/mosaic_res) + # maxres = math.sqrt((Sx*Sy)/(GPUmem*0.00008)) + # if maxres < 1: + # maxres = 1 + # img = cv2.resize(img, (int(Sx/maxres),int(Sy/maxres))) + img = img * 1.0 / 255 img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float() img_LR = img.unsqueeze(0) + + # image to device img_LR = img_LR.to(self.device) output = self.model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()