Spaces:
Configuration error
Configuration error
| import numpy as np | |
| import cv2 | |
| from PIL import Image | |
| import base64 | |
| import io | |
| from typing import List | |
| def encode_image_to_base64(img_array : np.ndarray): | |
| """Convert a numpy image array to base64-encoded PNG.""" | |
| # If float, convert to uint8 | |
| if np.issubdtype(img_array.dtype, np.floating): | |
| img_array = np.clip(img_array * 255, 0, 255).astype(np.uint8) | |
| img_pil = Image.fromarray(img_array) | |
| buf = io.BytesIO() | |
| img_pil.save(buf, format="PNG") | |
| return base64.b64encode(buf.getvalue()).decode('utf-8') | |
| def crop_to_mask_size(output, mask): | |
| """ | |
| Crops the output image to the bounding box of the mask. | |
| Parameters: | |
| output: The image to be cropped (HxWxC array). | |
| mask: The binary mask (HxW array). | |
| Returns: | |
| Cropped image (HxWxC array). | |
| """ | |
| if mask.ndim == 3 and mask.shape[0] == 1: | |
| mask = mask[0] # -> now (H, W) | |
| ys, xs = np.where(mask) | |
| if xs.size and ys.size: | |
| x0, x1 = xs.min(), xs.max() | |
| y0, y1 = ys.min(), ys.max() | |
| return output[y0:y1+1, x0:x1+1] | |
| return output | |
| def apply_mask(image, mask, background_mode="remove"): | |
| """ | |
| Applies the given mask to the image based on the specified background mode. | |
| Parameters: | |
| image: The input image (HxWx3 array). | |
| mask: The binary mask (HxW array). | |
| background_mode: The mode for background processing ("extreme_blur", "highlight", or "remove"). | |
| Returns: | |
| Processed image (HxWx3 uint8 array). | |
| """ | |
| # If the mask has a leading batch dimension, squeeze it | |
| if mask.ndim == 3 and mask.shape[0] == 1: | |
| mask = mask[0] # -> now (H, W) | |
| # Expand the mask into a 3-channel float mask of shape (H, W, 3) | |
| mask_3c = np.repeat(mask[..., None], 3, axis=2).astype(np.float32) | |
| img_f = image.astype(np.float32) | |
| one_c = 1.0 - mask_3c | |
| if background_mode == "extreme_blur": | |
| blurred = cv2.GaussianBlur(image, (101, 101), 0).astype(np.float32) | |
| output_f = img_f * mask_3c + blurred * one_c | |
| elif background_mode == "highlight": | |
| alpha = 0.5 | |
| overlay_color = np.array([255, 0, 0], dtype=np.float32) # pure red | |
| output_f = img_f.copy() | |
| output_f[mask] = (1 - alpha) * img_f[mask] + alpha * overlay_color | |
| else: # "remove" | |
| white = np.full_like(img_f, 255, dtype=np.float32) | |
| output_f = img_f * mask_3c + white * one_c | |
| return output_f.astype(np.uint8) | |
| def resize_image(image, max_width=800): | |
| """ | |
| Resizes a numpy array image (RGB) to a maximum width of 800px, preserving aspect ratio. | |
| """ | |
| if image is None: | |
| return None | |
| pil_img = Image.fromarray(image) | |
| width, height = pil_img.size | |
| if width > max_width: | |
| new_height = int(height * (max_width / width)) | |
| resized_img = pil_img.resize((max_width, new_height), Image.LANCZOS) | |
| return np.array(resized_img) | |
| else: | |
| return image |