segmentation/utils.py

"""
Utility functions for segmentation visualization and I/O.
"""

import numpy as np
import cv2
from PIL import Image
from typing import Tuple


def visualize_mask(
    image: Image.Image,
    mask: np.ndarray,
    alpha: float = 0.5,
    color: Tuple[int, int, int] = (255, 0, 0)
) -> Image.Image:
    """
    Overlay segmentation mask on image.
    
    Args:
        image: PIL Image
        mask: Binary mask (H, W)
        alpha: Transparency (0-1)
        color: RGB color tuple for mask
    
    Returns:
        Image with mask overlay
    """
    # Convert to numpy
    img_array = np.array(image)
    
    # Create colored mask
    colored_mask = np.zeros_like(img_array)
    colored_mask[mask > 0.5] = color
    
    # Blend
    result = cv2.addWeighted(img_array, 1.0, colored_mask, alpha, 0)
    
    return Image.fromarray(result)


def save_mask(mask: np.ndarray, output_path: str):
    """
    Save mask as image (white for ROI, black for background).
    
    Args:
        mask: Binary mask array (H, W)
        output_path: Path to save mask image
    """
    mask_img = (mask * 255).astype(np.uint8)
    Image.fromarray(mask_img).save(output_path)


def load_mask(mask_path: str) -> np.ndarray:
    """
    Load mask from image file.
    
    Args:
        mask_path: Path to mask image
    
    Returns:
        Binary mask as numpy array (H, W) with values 0 or 1
    """
    mask_img = Image.open(mask_path).convert('L')
    mask = np.array(mask_img).astype(np.float32) / 255.0
    return mask


def calculate_roi_stats(mask: np.ndarray) -> dict:
    """
    Calculate statistics about ROI coverage.
    
    Args:
        mask: Binary mask (H, W)
    
    Returns:
        Dictionary with statistics:
        - roi_pixels: Number of ROI pixels
        - total_pixels: Total number of pixels
        - roi_percentage: Percentage of image covered by ROI
    """
    roi_pixels = int(np.sum(mask > 0.5))
    total_pixels = int(mask.size)
    roi_percentage = (roi_pixels / total_pixels) * 100
    
    return {
        'roi_pixels': roi_pixels,
        'total_pixels': total_pixels,
        'roi_percentage': roi_percentage
    }