| import torch |
| import numpy as np |
| from PIL import Image |
| import gc |
| import numpy as np |
| import numpy as np |
| from PIL import Image |
| from scripts.refine_lr_to_sr import run_sr_fast |
|
|
| GRADIO_CACHE = "/tmp/gradio/" |
|
|
| def clean_up(): |
| torch.cuda.empty_cache() |
| gc.collect() |
|
|
| def remove_color(arr): |
| if arr.shape[-1] == 4: |
| arr = arr[..., :3] |
| |
| base = arr[0, 0] |
| diffs = np.abs(arr.astype(np.int32) - base.astype(np.int32)).sum(axis=-1) |
| alpha = (diffs <= 80) |
| |
| arr[alpha] = 255 |
| alpha = ~alpha |
| arr = np.concatenate([arr, alpha[..., None].astype(np.int32) * 255], axis=-1) |
| return arr |
|
|
| def simple_remove(imgs, run_sr=True): |
| """Only works for normal""" |
| if not isinstance(imgs, list): |
| imgs = [imgs] |
| single_input = True |
| else: |
| single_input = False |
| if run_sr: |
| imgs = run_sr_fast(imgs) |
| rets = [] |
| for img in imgs: |
| arr = np.array(img) |
| arr = remove_color(arr) |
| rets.append(Image.fromarray(arr.astype(np.uint8))) |
| if single_input: |
| return rets[0] |
| return rets |
|
|
| def rgba_to_rgb(rgba: Image.Image, bkgd="WHITE"): |
| new_image = Image.new("RGBA", rgba.size, bkgd) |
| new_image.paste(rgba, (0, 0), rgba) |
| new_image = new_image.convert('RGB') |
| return new_image |
|
|
| def change_rgba_bg(rgba: Image.Image, bkgd="WHITE"): |
| rgb_white = rgba_to_rgb(rgba, bkgd) |
| new_rgba = Image.fromarray(np.concatenate([np.array(rgb_white), np.array(rgba)[:, :, 3:4]], axis=-1)) |
| return new_rgba |
|
|
| def split_image(image, rows=None, cols=None): |
| """ |
| inverse function of make_image_grid |
| """ |
| |
| if rows is None and cols is None: |
| |
| rows = 1 |
| cols = image.size[0] // image.size[1] |
| assert cols * image.size[1] == image.size[0] |
| subimg_size = image.size[1] |
| elif rows is None: |
| subimg_size = image.size[0] // cols |
| rows = image.size[1] // subimg_size |
| assert rows * subimg_size == image.size[1] |
| elif cols is None: |
| subimg_size = image.size[1] // rows |
| cols = image.size[0] // subimg_size |
| assert cols * subimg_size == image.size[0] |
| else: |
| subimg_size = image.size[1] // rows |
| assert cols * subimg_size == image.size[0] |
| subimgs = [] |
| for i in range(rows): |
| for j in range(cols): |
| subimg = image.crop((j*subimg_size, i*subimg_size, (j+1)*subimg_size, (i+1)*subimg_size)) |
| subimgs.append(subimg) |
| return subimgs |
|
|
| def make_image_grid(images, rows=None, cols=None, resize=None): |
| if rows is None and cols is None: |
| rows = 1 |
| cols = len(images) |
| if rows is None: |
| rows = len(images) // cols |
| if len(images) % cols != 0: |
| rows += 1 |
| if cols is None: |
| cols = len(images) // rows |
| if len(images) % rows != 0: |
| cols += 1 |
| total_imgs = rows * cols |
| if total_imgs > len(images): |
| images += [Image.new(images[0].mode, images[0].size) for _ in range(total_imgs - len(images))] |
| |
| if resize is not None: |
| images = [img.resize((resize, resize)) for img in images] |
|
|
| w, h = images[0].size |
| grid = Image.new(images[0].mode, size=(cols * w, rows * h)) |
|
|
| for i, img in enumerate(images): |
| grid.paste(img, box=(i % cols * w, i // cols * h)) |
| return grid |
|
|
|
|
