Spaces:
Configuration error
Configuration error
File size: 2,924 Bytes
24f3fb6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
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 |