| import cv2 |
| import numpy as np |
| import torch |
| import os |
| import tempfile |
| from app.models.ESRGAN import RRDBNet_arch as arch |
|
|
| _MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models', 'RRDB_ESRGAN_x4.pth') |
| _device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| _model = None |
|
|
|
|
| def _get_model(): |
| global _model |
| if _model is None: |
| _model = arch.RRDBNet(3, 3, 64, 23, gc=32) |
| _model.load_state_dict(torch.load(_MODEL_PATH, map_location=_device), strict=True) |
| _model.eval() |
| _model = _model.to(_device) |
| return _model |
|
|
|
|
| def upscale_image(image_filepath): |
| model = _get_model() |
|
|
| print('Up-scaling with device: {}...'.format(_device)) |
|
|
| image = cv2.imread(image_filepath, cv2.IMREAD_COLOR) |
| image = image * 1.0 / 255 |
| image = torch.from_numpy(np.transpose(image[:, :, [2, 1, 0]], (2, 0, 1))).float() |
| image_low_res = image.unsqueeze(0).to(_device) |
|
|
| with torch.no_grad(): |
| image_high_res = model(image_low_res).data.squeeze().float().cpu().clamp_(0, 1).numpy() |
| image_high_res = np.transpose(image_high_res[[2, 1, 0], :, :], (1, 2, 0)) |
| image_high_res = (image_high_res * 255.0).round() |
|
|
| _, output_filepath = tempfile.mkstemp(suffix='.png') |
| cv2.imwrite(output_filepath, image_high_res) |
| print("image saved as:", output_filepath) |
|
|
| return output_filepath |
|
|
|
|
|
|
|
|
