masks.py

from typing import Tuple

import cv2
import numpy as np


def generate_naive_mask(
    rgb_image: np.ndarray,
    min_percentage: float = 1.0,
    max_percentage: float = 90.0
) -> Tuple[np.ndarray, np.ndarray, float, bool]:
    """Generate a naive foreground mask using brightness + Otsu thresholding.

    Returns:
        mask_bool: Boolean mask (H, W)
        debug_mask: uint8 mask for visualization (H, W)
        mask_percentage: % of pixels active in mask_bool
        fallback_full_image: True if the mask was replaced by full-image mask
    """
    gray = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2GRAY)
    _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

    kernel = np.ones((5, 5), np.uint8)
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)  # Fill holes
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)   # Remove noise

    debug_mask = mask.copy()
    mask_bool = mask.astype(bool)
    mask_percentage = (mask_bool.sum() / mask_bool.size) * 100

    fallback_full_image = False
    if mask_percentage < min_percentage or mask_percentage > max_percentage:
        fallback_full_image = True
        mask_bool = np.ones((rgb_image.shape[0], rgb_image.shape[1]), dtype=bool)
        debug_mask = np.ones((rgb_image.shape[0], rgb_image.shape[1]), dtype=np.uint8) * 255

    return mask_bool, debug_mask, mask_percentage, fallback_full_image