Skip to content

Commit

Permalink
Removed prints
Browse files Browse the repository at this point in the history
  • Loading branch information
natethegreate committed May 21, 2020
1 parent b742456 commit cc6e549
Showing 1 changed file with 13 additions and 12 deletions.
25 changes: 13 additions & 12 deletions ColabESRGAN/test.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import numpy as np
import torch
import architecture
import math

# ESRGAN class allows abstraction of warmup and inference.
class esrgan():
Expand All @@ -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()
Expand Down

0 comments on commit cc6e549

Please sign in to comment.