add image tiling from Real-ESRGAN to save memory

ESRGANer.py

@@ -0,0 +1,152 @@
+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)
+
+        # 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]
+
+    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
+    
+    @torch.no_grad()
+    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

app.py

@@ -3,8 +3,6 @@ import util
 import process_image
 from run_cmd import run_cmd
 
-run_cmd("pip install split-image")
-
 is_colab = util.is_google_colab()
 
 css = '''
@@ -37,7 +35,11 @@ with gr.Blocks(title=title, css=css) as demo:
     # {title}
     This space uses old ESRGAN architecture to upscale images, using models made by the community.
 
-    Once upscaled, click or tap the download button under the image to download it.
+    Once the photo upscaled, click or tap the **download button** under the image to download it. **The preview image is not the upscaled one**
+
+    I'll add more models after optimizing to size of the output image, right now it could be quite big.
+    
+    **Colab coming soon™**
     """)
 
     with gr.Box():

inference.py

@@ -4,8 +4,8 @@ import cv2
 import numpy as np
 import torch
 import architecture as arch
-from split_image import split_image
 from run_cmd import run_cmd
+from ESRGANer import ESRGANer
 
 def is_cuda():
     if torch.cuda.is_available():
@@ -39,37 +39,17 @@ for k, v in model.named_parameters():
     v.requires_grad = False
 model = model.to(device)
 
-base_path = os.path.dirname(img_path)
-image_name = os.path.basename(img_path)
+# Read image
+img = cv2.imread(img_path, cv2.IMREAD_COLOR)
+img = img * 1.0 / 255
+img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
+img_LR = img.unsqueeze(0)
+img_LR = img_LR.to(device)
 
-# Split image
-run_cmd(f"split-image {img_path} 7 7 --output-dir {base_path} --quiet")
+upsampler = ESRGANer(model=model)
+output = upsampler.enhance(img_LR)
 
-for root, dirs, files in os.walk(base_path, topdown=True):
-    for x, name in enumerate(files):
-        file_path = os.path.join(root, name)
-
-        if file_path == img_path:
-            continue
-
-        # Read image
-        img = cv2.imread(file_path, cv2.IMREAD_COLOR)
-        img = img * 1.0 / 255
-        img = torch.from_numpy(np.transpose(img[:, :, [2, 1, 0]], (2, 0, 1))).float()
-        img_LR = img.unsqueeze(0)
-        img_LR = img_LR.to(device)
-
-        print(f"Start upscaling tile {x}...")
-        with torch.no_grad():
-            output = model(img_LR).data.squeeze().float().cpu().clamp_(0, 1).numpy()
-        output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
-        output = (output * 255.0).round()
-        print(f"Finished upscaling tile {x}, saving tile.")
-        cv2.imwrite(file_path, output)
-
-# Join all tiles
-run_cmd(f"cd {base_path} && split-image {image_name} 7 7 -r --quiet")
-
-# Open image and save as png with the ouput name
-img_out = cv2.imread(img_path);
-cv2.imwrite(output_dir, img_out, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])
+output = output.squeeze().float().cpu().clamp_(0, 1).numpy()
+output = np.transpose(output[[2, 1, 0], :, :], (1, 2, 0))
+output = (output * 255.0).round()
+cv2.imwrite(output_dir, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])

inference_manga_v2.py

@@ -4,8 +4,7 @@ import cv2
 import numpy as np
 import torch
 import architecture as arch
-from split_image import split_image
-from run_cmd import run_cmd
+from ESRGANer import ESRGANer
 
 def is_cuda():
     if torch.cuda.is_available():
@@ -33,38 +32,17 @@ for k, v in model.named_parameters():
     v.requires_grad = False
 model = model.to(device)
 
-base_path = os.path.dirname(img_path)
-image_name = os.path.basename(img_path)
+# Read image
+img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+img = img * 1.0 / 255
+img = torch.from_numpy(img[np.newaxis, :, :]).float()
+img_LR = img.unsqueeze(0)
+img_LR = img_LR.to(device)
 
-# Split image
-run_cmd(f"split-image {img_path} 7 7 --output-dir {base_path} --quiet")
-
-for root, dirs, files in os.walk(base_path, topdown=True):
-    for x, name in enumerate(files):
-        file_path = os.path.join(root, name)
-
-        if file_path == img_path:
-            continue
-
-        # Read image
-        img = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
-        img = img * 1.0 / 255
-        img = torch.from_numpy(img[np.newaxis, :, :]).float()
-        img_LR = img.unsqueeze(0)
-        img_LR = img_LR.to(device)
-
-        print(f"Start upscaling tile {x}...")
-        with torch.no_grad():
-            output = model(img_LR).squeeze(dim=0).float().cpu().clamp_(0, 1).numpy()
-        output = np.transpose(output, (1, 2, 0))
-        output = (output * 255.0).round()
-        print(f"Finished upscaling tile {x}, saving tile.")
-        cv2.imwrite(file_path, output)
-
-# Join all tiles
-run_cmd(f"cd {base_path} && split-image {image_name} 7 7 -r --quiet")
-
-# Open image and save as png with the ouput name
-img_out = cv2.imread(img_path);
-cv2.imwrite(output_dir, img_out, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])
+upsampler = ESRGANer(model=model)
+output = upsampler.enhance(img_LR)
 
+output = output.squeeze(dim=0).float().cpu().clamp_(0, 1).numpy()
+output = np.transpose(output, (1, 2, 0))
+output = (output * 255.0).round()
+cv2.imwrite(output_dir, output, [int(cv2.IMWRITE_PNG_COMPRESSION), 5])

process_image.py

@@ -16,10 +16,9 @@ def inference(img, size, type):
     OUTPUT_DIR = os.path.join(temp_path, f"output_image{str(_id)}")
     img_in_path = os.path.join(INPUT_DIR, "input.jpg")
     img_out_path = os.path.join(OUTPUT_DIR, f"output_{size}.png")
-    run_cmd(f"rm -rf {INPUT_DIR}")
-    run_cmd(f"rm -rf {OUTPUT_DIR}")
     run_cmd(f"mkdir {INPUT_DIR}")
     run_cmd(f"mkdir {OUTPUT_DIR}")
+    
     img.save(img_in_path, "PNG")
 
     if type == "Manga":
@@ -32,9 +31,6 @@ def inference(img, size, type):
     if size == "x2":
         img_out = img_out.resize((img_out.width // 2, img_out.height // 2), resample=Image.BICUBIC)
 
-    #img_out.save(img_out_path, optimize=True) # Add more optimizations
-    #img_out = Image.open(img_out_path)
-
     # Remove input and output image
     run_cmd(f"rm -rf {INPUT_DIR}")