Spaces:
Running
Running
| """Preprocessing methods""" | |
| import logging | |
| from typing import List, Tuple | |
| import numpy as np | |
| from PIL import Image, ImageFilter | |
| import streamlit as st | |
| from config import COLOR_RGB, WIDTH, HEIGHT | |
| # from enhance_config import ENHANCE_SETTINGS | |
| LOGGING = logging.getLogger(__name__) | |
| def preprocess_seg_mask(canvas_seg, real_seg: Image.Image = None) -> Tuple[np.ndarray, np.ndarray]: | |
| """Preprocess the segmentation mask. | |
| Args: | |
| canvas_seg: segmentation canvas | |
| real_seg (Image.Image, optional): segmentation mask. Defaults to None. | |
| Returns: | |
| Tuple[np.ndarray, np.ndarray]: segmentation mask, segmentation mask with overlay | |
| """ | |
| # get unique colors in the segmentation | |
| image_seg = canvas_seg.image_data.copy()[:, :, :3] | |
| # average the colors of the segmentation masks | |
| average_color = np.mean(image_seg, axis=(2)) | |
| mask = average_color[:, :] > 0 | |
| if mask.sum() > 0: | |
| mask = mask * 1 | |
| unique_colors = np.unique(image_seg.reshape(-1, image_seg.shape[-1]), axis=0) | |
| unique_colors = [tuple(color) for color in unique_colors] | |
| unique_colors = [color for color in unique_colors if np.sum( | |
| np.all(image_seg == color, axis=-1)) > 100] | |
| unique_colors_exact = [color for color in unique_colors if color in COLOR_RGB] | |
| if real_seg is not None: | |
| overlay_seg = np.array(real_seg) | |
| unique_colors = np.unique(overlay_seg.reshape(-1, overlay_seg.shape[-1]), axis=0) | |
| unique_colors = [tuple(color) for color in unique_colors] | |
| for color in unique_colors_exact: | |
| if color != (255, 255, 255) and color != (0, 0, 0): | |
| overlay_seg[np.all(image_seg == color, axis=-1)] = color | |
| image_seg = overlay_seg | |
| return mask, image_seg | |
| def get_mask(image_mask: np.ndarray) -> np.ndarray: | |
| """Get the mask from the segmentation mask. | |
| Args: | |
| image_mask (np.ndarray): segmentation mask | |
| Returns: | |
| np.ndarray: mask | |
| """ | |
| # average the colors of the segmentation masks | |
| average_color = np.mean(image_mask, axis=(2)) | |
| mask = average_color[:, :] > 0 | |
| if mask.sum() > 0: | |
| mask = mask * 1 | |
| return mask | |
| def get_image() -> np.ndarray: | |
| """Get the image from the session state. | |
| Returns: | |
| np.ndarray: image | |
| """ | |
| if 'initial_image' in st.session_state and st.session_state['initial_image'] is not None: | |
| initial_image = st.session_state['initial_image'] | |
| if isinstance(initial_image, Image.Image): | |
| return np.array(initial_image.resize((WIDTH, HEIGHT))) | |
| else: | |
| return np.array(Image.fromarray(initial_image).resize((WIDTH, HEIGHT))) | |
| else: | |
| return None | |
| # def make_enhance_config(segmentation, objects=None): | |
| """Make the enhance config for the segmentation image. | |
| """ | |
| info = ENHANCE_SETTINGS[objects] | |
| segmentation = np.array(segmentation) | |
| if 'replace' in info: | |
| replace_color = info['replace'] | |
| mask = np.zeros(segmentation.shape) | |
| for color in info['colors']: | |
| mask[np.all(segmentation == color, axis=-1)] = [1, 1, 1] | |
| segmentation[np.all(segmentation == color, axis=-1)] = replace_color | |
| if info['inverse'] is False: | |
| mask = np.zeros(segmentation.shape) | |
| for color in info['colors']: | |
| mask[np.all(segmentation == color, axis=-1)] = [1, 1, 1] | |
| else: | |
| mask = np.ones(segmentation.shape) | |
| for color in info['colors']: | |
| mask[np.all(segmentation == color, axis=-1)] = [0, 0, 0] | |
| st.session_state['positive_prompt'] = info['positive_prompt'] | |
| st.session_state['negative_prompt'] = info['negative_prompt'] | |
| if info['inpainting'] is True: | |
| mask = mask.astype(np.uint8) | |
| mask = Image.fromarray(mask) | |
| mask = mask.filter(ImageFilter.GaussianBlur(radius=13)) | |
| mask = mask.filter(ImageFilter.MaxFilter(size=9)) | |
| mask = np.array(mask) | |
| mask[mask < 0.1] = 0 | |
| mask[mask >= 0.1] = 1 | |
| mask = mask.astype(np.uint8) | |
| conditioning = dict( | |
| mask_image=mask, | |
| positive_prompt=info['positive_prompt'], | |
| negative_prompt=info['negative_prompt'], | |
| ) | |
| else: | |
| conditioning = dict( | |
| mask_image=mask, | |
| controlnet_conditioning_image=segmentation, | |
| positive_prompt=info['positive_prompt'], | |
| negative_prompt=info['negative_prompt'], | |
| strength=info['strength'] | |
| ) | |
| return conditioning, info['inpainting'] |