Spaces:
Runtime error
Runtime error
| from PIL import Image, ImageChops | |
| import numpy as np | |
| import cv2 as cv | |
| import math | |
| import torch | |
| from torch.nn import functional as F | |
| """ | |
| Borrowed and adapted from https://github.com/xinntao/Real-ESRGAN/blob/master/realesrgan/utils.py | |
| Thank you xinntao! | |
| """ | |
| class ESRGANer(): | |
| """A helper class for upsampling images with ESRGAN. | |
| Args: | |
| scale (int): Upsampling scale factor used in the networks. It is usually 2 or 4. | |
| model (nn.Module): The defined network. Default: None. | |
| tile (int): As too large images result in the out of GPU memory issue, so this tile option will first crop | |
| input images into tiles, and then process each of them. Finally, they will be merged into one image. | |
| 0 denotes for do not use tile. Default: 500. | |
| tile_pad (int): The pad size for each tile, to remove border artifacts. Default: 10. | |
| pre_pad (int): Pad the input images to avoid border artifacts. Default: 10. | |
| """ | |
| def __init__(self, | |
| scale=4, | |
| model=None, | |
| tile=300, | |
| tile_pad=10, | |
| pre_pad=10 | |
| ): | |
| self.scale = scale | |
| self.tile_size = tile | |
| self.tile_pad = tile_pad | |
| self.pre_pad = pre_pad | |
| self.mod_scale = None | |
| self.model = model | |
| def pre_process(self, img): | |
| """Pre-process, such as pre-pad and mod pad, so that the images can be divisible | |
| """ | |
| self.img = img | |
| # pre_pad | |
| if self.pre_pad != 0: | |
| self.img = F.pad(self.img, (0, self.pre_pad, 0, self.pre_pad), 'reflect') | |
| # mod pad for divisible borders | |
| if self.scale == 2: | |
| self.mod_scale = 2 | |
| elif self.scale == 1: | |
| self.mod_scale = 4 | |
| if self.mod_scale is not None: | |
| self.mod_pad_h, self.mod_pad_w = 0, 0 | |
| _, _, h, w = self.img.size() | |
| if (h % self.mod_scale != 0): | |
| self.mod_pad_h = (self.mod_scale - h % self.mod_scale) | |
| if (w % self.mod_scale != 0): | |
| self.mod_pad_w = (self.mod_scale - w % self.mod_scale) | |
| self.img = F.pad(self.img, (0, self.mod_pad_w, 0, self.mod_pad_h), 'reflect') | |
| def process(self): | |
| # model inference | |
| self.output = self.model(self.img) | |
| def tile_process(self): | |
| """It will first crop input images to tiles, and then process each tile. | |
| Finally, all the processed tiles are merged into one images. | |
| Modified from: https://github.com/ata4/esrgan-launcher | |
| """ | |
| batch, channel, height, width = self.img.shape | |
| output_height = height * self.scale | |
| output_width = width * self.scale | |
| output_shape = (batch, channel, output_height, output_width) | |
| # start with black image | |
| self.output = self.img.new_zeros(output_shape) | |
| tiles_x = math.ceil(width / self.tile_size) | |
| tiles_y = math.ceil(height / self.tile_size) | |
| print("Image processing started...") | |
| # loop over all tiles | |
| for y in range(tiles_y): | |
| for x in range(tiles_x): | |
| # extract tile from input image | |
| ofs_x = x * self.tile_size | |
| ofs_y = y * self.tile_size | |
| # input tile area on total image | |
| input_start_x = ofs_x | |
| input_end_x = min(ofs_x + self.tile_size, width) | |
| input_start_y = ofs_y | |
| input_end_y = min(ofs_y + self.tile_size, height) | |
| # input tile area on total image with padding | |
| input_start_x_pad = max(input_start_x - self.tile_pad, 0) | |
| input_end_x_pad = min(input_end_x + self.tile_pad, width) | |
| input_start_y_pad = max(input_start_y - self.tile_pad, 0) | |
| input_end_y_pad = min(input_end_y + self.tile_pad, height) | |
| # input tile dimensions | |
| input_tile_width = input_end_x - input_start_x | |
| input_tile_height = input_end_y - input_start_y | |
| tile_idx = y * tiles_x + x + 1 | |
| input_tile = self.img[:, :, input_start_y_pad:input_end_y_pad, input_start_x_pad:input_end_x_pad] | |
| # upscale tile | |
| try: | |
| with torch.no_grad(): | |
| output_tile = self.model(input_tile) | |
| except RuntimeError as error: | |
| print('Error', error) | |
| print(f'Processing tile {tile_idx}/{tiles_x * tiles_y}') | |
| # output tile area on total image | |
| output_start_x = input_start_x * self.scale | |
| output_end_x = input_end_x * self.scale | |
| output_start_y = input_start_y * self.scale | |
| output_end_y = input_end_y * self.scale | |
| # output tile area without padding | |
| output_start_x_tile = (input_start_x - input_start_x_pad) * self.scale | |
| output_end_x_tile = output_start_x_tile + input_tile_width * self.scale | |
| output_start_y_tile = (input_start_y - input_start_y_pad) * self.scale | |
| output_end_y_tile = output_start_y_tile + input_tile_height * self.scale | |
| # put tile into output image | |
| self.output[:, :, output_start_y:output_end_y, | |
| output_start_x:output_end_x] = output_tile[:, :, output_start_y_tile:output_end_y_tile, | |
| output_start_x_tile:output_end_x_tile] | |
| print('All tiles processed, saving output image!') | |
| def post_process(self): | |
| # remove extra pad | |
| if self.mod_scale is not None: | |
| _, _, h, w = self.output.size() | |
| self.output = self.output[:, :, 0:h - self.mod_pad_h * self.scale, 0:w - self.mod_pad_w * self.scale] | |
| # remove prepad | |
| if self.pre_pad != 0: | |
| _, _, h, w = self.output.size() | |
| self.output = self.output[:, :, 0:h - self.pre_pad * self.scale, 0:w - self.pre_pad * self.scale] | |
| return self.output | |
| def enhance(self, img): | |
| self.pre_process(img) | |
| if self.tile_size > 0: | |
| self.tile_process() | |
| else: | |
| self.process() | |
| output_img = self.post_process() | |
| return output_img |