subject_mask.py

import numpy as np
import cv2

def depth_foreground_mask(depth01: np.ndarray, strength: float = 1.0) -> np.ndarray:
    """
    Fast foreground mask from depth.
    strength: 0..2 (higher = more aggressive background removal)
    Returns float mask 0..1
    """
    d = depth01.astype(np.float32)
    # threshold using percentile
    p = 30 + int(10 * strength)  # 30..50
    thr = np.percentile(d, p)
    mask = (d > thr).astype(np.float32)

    # morphology cleanup
    k = int(9 + 6 * strength)  # 9..21
    k = k if k % 2 == 1 else k + 1
    ker = np.ones((k, k), np.uint8)
    mask_u8 = (mask * 255).astype(np.uint8)
    mask_u8 = cv2.morphologyEx(mask_u8, cv2.MORPH_OPEN, ker)
    mask_u8 = cv2.morphologyEx(mask_u8, cv2.MORPH_CLOSE, ker)

    # feather edges
    mask_f = mask_u8.astype(np.float32) / 255.0
    mask_f = cv2.GaussianBlur(mask_f, (k, k), 0)
    return np.clip(mask_f, 0.0, 1.0)