File size: 4,130 Bytes
0ed69a6 |
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 |
import torch
import base64
import gradio as gr
import numpy as np
from PIL import Image
from io import BytesIO
MAX_COLORS = 12
def create_binary_matrix(img_arr, target_color):
mask = np.all(img_arr == target_color, axis=-1)
binary_matrix = mask.astype(int)
return binary_matrix
def preprocess_mask(mask_, h, w, device):
mask = np.array(mask_)
mask = mask.astype(np.float32)
mask = mask[None, None]
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
mask = torch.from_numpy(mask).to(device)
mask = torch.nn.functional.interpolate(mask, size=(h, w), mode='nearest')
return mask
def process_sketch(canvas_data):
binary_matrixes = []
base64_img = canvas_data['image']
image_data = base64.b64decode(base64_img.split(',')[1])
image = Image.open(BytesIO(image_data)).convert("RGB")
im2arr = np.array(image)
colors = [tuple(map(int, rgb[4:-1].split(','))) for rgb in canvas_data['colors']]
colors_fixed = []
r, g, b = 255, 255, 255
binary_matrix = create_binary_matrix(im2arr, (r,g,b))
binary_matrixes.append(binary_matrix)
binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))
colored_map = binary_matrix_*(r,g,b) + (1-binary_matrix_)*(50,50,50)
colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))
for color in colors:
r, g, b = color
if any(c != 255 for c in (r, g, b)):
binary_matrix = create_binary_matrix(im2arr, (r,g,b))
binary_matrixes.append(binary_matrix)
binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))
colored_map = binary_matrix_*(r,g,b) + (1-binary_matrix_)*(50,50,50)
colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))
visibilities = []
colors = []
for n in range(MAX_COLORS):
visibilities.append(gr.update(visible=False))
colors.append(gr.update())
for n in range(len(colors_fixed)):
visibilities[n] = gr.update(visible=True)
colors[n] = colors_fixed[n]
return [gr.update(visible=True), binary_matrixes, *visibilities, *colors]
def process_prompts(binary_matrixes, *seg_prompts):
return [gr.update(visible=True), gr.update(value=' , '.join(seg_prompts[:len(binary_matrixes)]))]
def process_example(layout_path, all_prompts, seed_):
all_prompts = all_prompts.split('***')
binary_matrixes = []
colors_fixed = []
im2arr = np.array(Image.open(layout_path))[:,:,:3]
unique, counts = np.unique(np.reshape(im2arr,(-1,3)), axis=0, return_counts=True)
sorted_idx = np.argsort(-counts)
binary_matrix = create_binary_matrix(im2arr, (0,0,0))
binary_matrixes.append(binary_matrix)
binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))
colored_map = binary_matrix_*(255,255,255) + (1-binary_matrix_)*(50,50,50)
colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))
for i in range(len(all_prompts)-1):
r, g, b = unique[sorted_idx[i]]
if any(c != 255 for c in (r, g, b)) and any(c != 0 for c in (r, g, b)):
binary_matrix = create_binary_matrix(im2arr, (r,g,b))
binary_matrixes.append(binary_matrix)
binary_matrix_ = np.repeat(np.expand_dims(binary_matrix, axis=(-1)), 3, axis=(-1))
colored_map = binary_matrix_*(r,g,b) + (1-binary_matrix_)*(50,50,50)
colors_fixed.append(gr.update(value=colored_map.astype(np.uint8)))
visibilities = []
colors = []
prompts = []
for n in range(MAX_COLORS):
visibilities.append(gr.update(visible=False))
colors.append(gr.update())
prompts.append(gr.update())
for n in range(len(colors_fixed)):
visibilities[n] = gr.update(visible=True)
colors[n] = colors_fixed[n]
prompts[n] = all_prompts[n+1]
return [gr.update(visible=True), binary_matrixes, *visibilities, *colors, *prompts,
gr.update(visible=True), gr.update(value=all_prompts[0]), int(seed_)]
|