| 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_)] | |