Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,360 Bytes
37aeb5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
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]
# calc diffs
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
"""
# image is in square
if rows is None and cols is None:
# image.size [W, H]
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
|