Spaces:
Sleeping
Sleeping
| """Preprocessing methods""" | |
| import logging | |
| from typing import List, Tuple | |
| import numpy as np | |
| from PIL import Image, ImageFilter | |
| from config import COLOR_RGB | |
| # 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 | |