67 lines
1.9 KiB
Python
67 lines
1.9 KiB
Python
"""
|
|
Image processing functions
|
|
"""
|
|
|
|
import numpy as np
|
|
import cv2
|
|
from typing import List
|
|
|
|
|
|
def resize_mask_to_image(
|
|
mask: np.ndarray,
|
|
target_height: int,
|
|
target_width: int
|
|
) -> np.ndarray:
|
|
"""
|
|
Resize mask to match image dimensions.
|
|
|
|
Args:
|
|
mask: Mask array (H, W) in range [0, 1]
|
|
target_height: Target height
|
|
target_width: Target width
|
|
|
|
Returns:
|
|
Resized mask array (target_height, target_width) in range [0, 1]
|
|
"""
|
|
if mask.shape[0] == target_height and mask.shape[1] == target_width:
|
|
return mask
|
|
|
|
# Use cv2.resize with INTER_LINEAR for smooth resizing
|
|
mask_resized = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_LINEAR)
|
|
# Ensure values stay in [0, 1] range
|
|
mask_resized = np.clip(mask_resized, 0.0, 1.0)
|
|
|
|
return mask_resized
|
|
|
|
|
|
def apply_mask_overlay(
|
|
image: np.ndarray,
|
|
mask: np.ndarray,
|
|
background_color: List[float]
|
|
) -> np.ndarray:
|
|
"""
|
|
Apply mask overlay to image with background color.
|
|
Where mask is 1.0, use background_color; where mask is 0.0, use original image.
|
|
|
|
Args:
|
|
image: Image array (H, W, 3) in range [0, 1]
|
|
mask: Mask array (H, W) in range [0, 1]
|
|
background_color: RGB color [r, g, b] in range [0, 1]
|
|
|
|
Returns:
|
|
Image with mask overlay applied (H, W, 3) in range [0, 1]
|
|
"""
|
|
# Ensure mask is 3D for broadcasting: (H, W) -> (H, W, 1)
|
|
mask_3d = np.expand_dims(mask, axis=-1)
|
|
|
|
# Expand background_color to match image shape: (3,) -> (1, 1, 3)
|
|
bg_color_array = np.array(background_color, dtype=np.float32).reshape(1, 1, 3)
|
|
|
|
# Apply overlay: image * (1 - mask) + background_color * mask
|
|
masked_image = image * (1.0 - mask_3d) + bg_color_array * mask_3d
|
|
|
|
# Ensure values stay in [0, 1] range
|
|
masked_image = np.clip(masked_image, 0.0, 1.0)
|
|
|
|
return masked_image
|