MWAI-Interior-Design / preprocessing.py
BertChristiaens's picture
Duplicate from BertChristiaens/controlnet-seg-backup
3d4d894
"""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']