Spaces:
Running
Running
| """ | |
| Source url: https://github.com/OPHoperHPO/image-background-remove-tool | |
| Author: Nikita Selin (OPHoperHPO)[https://github.com/OPHoperHPO]. | |
| License: Apache License 2.0 | |
| """ | |
| import cv2 | |
| import numpy as np | |
| from PIL import Image | |
| def prob_filter(mask: Image.Image, prob_threshold=231) -> Image.Image: | |
| """ | |
| Applies a filter to the mask by the probability of locating an object in the object area. | |
| Args: | |
| prob_threshold: Threshold of probability for mark area as background. | |
| mask: Predicted object mask | |
| Raises: | |
| ValueError if mask or trimap has wrong color mode | |
| Returns: | |
| Generated trimap for image. | |
| """ | |
| if mask.mode != "L": | |
| raise ValueError("Input mask has wrong color mode.") | |
| # noinspection PyTypeChecker | |
| mask_array = np.array(mask) | |
| mask_array[mask_array > prob_threshold] = 255 # Probability filter for mask | |
| mask_array[mask_array <= prob_threshold] = 0 | |
| return Image.fromarray(mask_array).convert("L") | |
| def prob_as_unknown_area( | |
| trimap: Image.Image, mask: Image.Image, prob_threshold=255 | |
| ) -> Image.Image: | |
| """ | |
| Marks any uncertainty in the seg mask as an unknown region. | |
| Args: | |
| prob_threshold: Threshold of probability for mark area as unknown. | |
| trimap: Generated trimap. | |
| mask: Predicted object mask | |
| Raises: | |
| ValueError if mask or trimap has wrong color mode | |
| Returns: | |
| Generated trimap for image. | |
| """ | |
| if mask.mode != "L" or trimap.mode != "L": | |
| raise ValueError("Input mask has wrong color mode.") | |
| # noinspection PyTypeChecker | |
| mask_array = np.array(mask) | |
| # noinspection PyTypeChecker | |
| trimap_array = np.array(trimap) | |
| trimap_array[np.logical_and(mask_array <= prob_threshold, mask_array > 0)] = 127 | |
| return Image.fromarray(trimap_array).convert("L") | |
| def post_erosion(trimap: Image.Image, erosion_iters=1) -> Image.Image: | |
| """ | |
| Performs erosion on the mask and marks the resulting area as an unknown region. | |
| Args: | |
| erosion_iters: The number of iterations of erosion that | |
| the object's mask will be subjected to before forming an unknown area | |
| trimap: Generated trimap. | |
| mask: Predicted object mask | |
| Returns: | |
| Generated trimap for image. | |
| """ | |
| if trimap.mode != "L": | |
| raise ValueError("Input mask has wrong color mode.") | |
| # noinspection PyTypeChecker | |
| trimap_array = np.array(trimap) | |
| if erosion_iters > 0: | |
| without_unknown_area = trimap_array.copy() | |
| without_unknown_area[without_unknown_area == 127] = 0 | |
| erosion_kernel = np.ones((3, 3), np.uint8) | |
| erode = cv2.erode( | |
| without_unknown_area, erosion_kernel, iterations=erosion_iters | |
| ) | |
| erode = np.where(erode == 0, 0, without_unknown_area) | |
| trimap_array[np.logical_and(erode == 0, without_unknown_area > 0)] = 127 | |
| erode = trimap_array.copy() | |
| else: | |
| erode = trimap_array.copy() | |
| return Image.fromarray(erode).convert("L") | |