01_tissue_detection.py

"""
01_tissue_detection.py — Tissue mask generation from multichannel IF images.

Generates a binary tissue mask by thresholding the image signal. By default,
uses a max-projection across all non-AF channels to capture tissue regions
regardless of which marker is expressed. Supports single-channel mode
(e.g., DAPI only) via --channels.

The mask is saved as a single-channel OME-TIFF. Tissue area (in pixels and
µm²) is printed and optionally saved as a parquet file.

Usage:
    python 01_tissue_detection.py --config config.toml
    python 01_tissue_detection.py --image_fullpath './image.ome.tiff'
"""

import os
import argparse

import numpy as np
import pandas as pd

import lib_sp as sp


# ============================================================================
#  Tissue detection core
# ============================================================================

def fn_threshold_image(
    arg_img,
    arg_downsample=4,
    arg_sigma=1.0,
    arg_threshold='otsu',
    arg_min_obj_size=2500,
    arg_max_hole_size=2500,
):
    """
    Threshold a 2D image to produce a binary tissue mask.

    Steps: downsample → Gaussian blur → threshold → remove small objects →
    fill small holes → upsample back to original size.

    Args:
        arg_img: 2D numpy array (or dask array) of the image.
        arg_downsample: Downsampling factor for speed (default: 4).
        arg_sigma: Gaussian blur sigma (default: 1.0).
        arg_threshold: Threshold value, or 'otsu' for automatic (default: 'otsu').
        arg_min_obj_size: Minimum object size in pixels at downsampled resolution (default: 2500).
        arg_max_hole_size: Maximum hole size to fill in pixels at downsampled resolution (default: 2500).

    Returns:
        arr_mask: Binary mask (uint16, 0 or 65535) at original resolution.
    """
    import cv2
    from skimage import filters
    from skimage.morphology import remove_small_objects, remove_small_holes

    # Downsample — .compute() handles dask arrays
    if (hasattr(arg_img, 'compute')):
        arr_img_downsampled = arg_img[::arg_downsample, ::arg_downsample].compute()
    else:
        arr_img_downsampled = arg_img[::arg_downsample, ::arg_downsample]

    # Gaussian blur
    arr_img_blurred = cv2.GaussianBlur(arr_img_downsampled, ksize=(-1, -1), sigmaX=arg_sigma)

    # Threshold
    if (arg_threshold == 'otsu'):
        flt_threshold = filters.threshold_otsu(arr_img_blurred)
        print(f'  Otsu threshold = {flt_threshold}')
    else:
        flt_threshold = arg_threshold
        print(f'  Fixed threshold = {flt_threshold}')

    arr_mask = arr_img_blurred > flt_threshold

    # Remove small objects and fill small holes
    try:
        # scikit-image >= 0.26
        arr_mask = remove_small_objects(arr_mask, max_size=arg_min_obj_size)
        arr_mask = remove_small_holes(arr_mask, max_size=arg_max_hole_size)
    except TypeError:
        # scikit-image < 0.26
        arr_mask = remove_small_objects(arr_mask, min_size=arg_min_obj_size)
        arr_mask = remove_small_holes(arr_mask, area_threshold=arg_max_hole_size)

    # Upsample back to original resolution
    if (arg_downsample != 1):
        tup_img_size_orig = (arg_img.shape[1], arg_img.shape[0])  # cv2.resize expects (width, height)
        arr_mask = cv2.resize(arr_mask.astype(np.uint8), tup_img_size_orig, interpolation=cv2.INTER_NEAREST)

    arr_mask = (arr_mask > 0).astype(np.uint16) * 65535

    return arr_mask


def fn_detect_tissue(
    arg_image_fullpath,
    arg_output_fullpath,
    arg_channels=None,
    arg_exclude_channels=None,
    arg_downsample=4,
    arg_sigma=1.0,
    arg_threshold='otsu',
    arg_min_obj_size=2500,
    arg_max_hole_size=2500,
    arg_overwrite=False,
):
    """
    Generate a binary tissue mask from a multichannel IF image.

    By default, max-projects across all channels except those in
    arg_exclude_channels (typically ['AF']). If arg_channels is specified,
    only those channels are used.

    Args:
        arg_image_fullpath: Path to the input OME-TIFF image.
        arg_output_fullpath: Path for the output tissue mask OME-TIFF.
        arg_channels: List of channel names to use. If None, uses all
            channels except arg_exclude_channels.
        arg_exclude_channels: List of channel names to exclude from
            max-projection (default: ['AF']). Ignored if arg_channels is set.
        arg_downsample: Downsampling factor (default: 4).
        arg_sigma: Gaussian blur sigma (default: 1.0).
        arg_threshold: Threshold value or 'otsu' (default: 'otsu').
        arg_min_obj_size: Min object size in pixels at downsampled resolution (default: 2500).
        arg_max_hole_size: Max hole size to fill in pixels at downsampled resolution (default: 2500).
        arg_overwrite: Overwrite existing output (default: False).

    Returns:
        arr_mask: Binary mask (uint16, 0 or 65535).
        flt_tissue_area_um2: Tissue area in µm² (None if pixel size unavailable).
    """
    if (os.path.exists(arg_output_fullpath) and not arg_overwrite):
        print(f'  Output already exists, skipping: {arg_output_fullpath}')
        print(f'  Use --overwrite to regenerate.')
        return None, None

    # Read image
    obj_img, dct_img_metadata = sp.fn_imread_bioio(arg_image_fullpath)
    lst_channel_names = dct_img_metadata['ChannelNames']
    print(f'  Image channels: {lst_channel_names}')

    # Determine which channels to use
    if (arg_channels is not None and arg_channels != ['all']):
        lst_use_channels = arg_channels
    else:
        # Use all channels except excluded ones
        lst_exclude = arg_exclude_channels if (arg_exclude_channels is not None) else ['AF']
        lst_use_channels = [ch for ch in lst_channel_names if ch not in lst_exclude]

    print(f'  Using channels for tissue detection: {lst_use_channels}')

    # Read channel data and max-project
    lst_channel_arrays = []
    for str_channel in lst_use_channels:
        if (str_channel not in lst_channel_names):
            print(f'  Warning: channel {str_channel} not found in image, skipping')
            continue
        int_channel_idx = lst_channel_names.index(str_channel)
        da_channel = obj_img.get_image_dask_data('YX', T=0, Z=0, C=int_channel_idx)
        lst_channel_arrays.append(da_channel)

    if (len(lst_channel_arrays) == 0):
        raise ValueError('No valid channels found for tissue detection')

    if (len(lst_channel_arrays) == 1):
        print(f'  Single channel mode: {lst_use_channels[0]}')
        da_img = lst_channel_arrays[0]
    else:
        print(f'  Max-projecting {len(lst_channel_arrays)} channels...')
        # Compute each channel at downsampled resolution, then max-project
        # This avoids loading all full-resolution channels into memory
        int_ds = arg_downsample
        arr_max = None
        for da_ch in lst_channel_arrays:
            if (hasattr(da_ch, 'compute')):
                arr_ch = da_ch[::int_ds, ::int_ds].compute()
            else:
                arr_ch = da_ch[::int_ds, ::int_ds]
            if (arr_max is None):
                arr_max = arr_ch.astype(np.float32)
            else:
                arr_max = np.maximum(arr_max, arr_ch.astype(np.float32))

        # Wrap in a simple object so fn_threshold_image can handle it
        # Since we already downsampled, set downsample=1
        arr_mask = fn_threshold_image(
            arg_img=arr_max,
            arg_downsample=1,
            arg_sigma=arg_sigma,
            arg_threshold=arg_threshold,
            arg_min_obj_size=arg_min_obj_size,
            arg_max_hole_size=arg_max_hole_size,
        )

        # Upsample back to original resolution
        if (int_ds != 1):
            import cv2
            tup_orig_size = (obj_img.shape[-1], obj_img.shape[-2])  # (width, height)
            arr_mask = cv2.resize(arr_mask.astype(np.uint8), tup_orig_size, interpolation=cv2.INTER_NEAREST)
            arr_mask = (arr_mask > 0).astype(np.uint16) * 65535

        # Skip the single-channel path below
        da_img = None

    # Single-channel path
    if (da_img is not None):
        arr_mask = fn_threshold_image(
            arg_img=da_img,
            arg_downsample=arg_downsample,
            arg_sigma=arg_sigma,
            arg_threshold=arg_threshold,
            arg_min_obj_size=arg_min_obj_size,
            arg_max_hole_size=arg_max_hole_size,
        )

    # Calculate tissue area
    int_tissue_pixels = int(np.sum(arr_mask > 0))
    int_total_pixels = arr_mask.shape[0] * arr_mask.shape[1]
    flt_tissue_fraction = int_tissue_pixels / int_total_pixels

    flt_tissue_area_um2 = None
    flt_res_x = dct_img_metadata.get('Resolution_X_um', None)
    flt_res_y = dct_img_metadata.get('Resolution_Y_um', None)

    if (flt_res_x is not None and flt_res_y is not None):
        flt_um2_per_px = flt_res_x * flt_res_y
        flt_tissue_area_um2 = int_tissue_pixels * flt_um2_per_px
        flt_tissue_area_mm2 = flt_tissue_area_um2 / 1e6
        print(f'  Pixel size: {flt_res_x:.4f} x {flt_res_y:.4f} µm')
        print(f'  Tissue area: {flt_tissue_area_mm2:.4f} mm² ({int_tissue_pixels:,} pixels, {flt_tissue_fraction:.1%} of image)')
    else:
        print(f'  Tissue pixels: {int_tissue_pixels:,} ({flt_tissue_fraction:.1%} of image)')
        print(f'  Warning: pixel size not available, cannot compute area in µm²')

    # Save mask
    os.makedirs(os.path.dirname(arg_output_fullpath), exist_ok=True)

    tup_pixel_size = (
        flt_res_x if (flt_res_x is not None) else 1.0,
        flt_res_y if (flt_res_y is not None) else 1.0,
    )

    sp.fnSavePyramidOMETIFF(
        arr_mask,
        arg_output_fullpath,
        ['tissue_mask'],
        tup_pixel_size,
        1,  # single pyramid level for mask
    )
    print(f'  Tissue mask saved: {arg_output_fullpath}')

    return arr_mask, flt_tissue_area_um2


# ============================================================================
#  CLI / TOML config
# ============================================================================

def fn_load_config_toml(arg_toml_fullpath: str) -> dict:
    """Load a TOML config file and return the [tissue_detection] section as a dict."""
    try:
        import tomllib  # Python 3.11+
    except ModuleNotFoundError:
        import tomli as tomllib  # Python 3.10 fallback
    with open(arg_toml_fullpath, 'rb') as f:
        dct_config = tomllib.load(f)
    return dct_config.get('tissue_detection', dct_config)


def fn_parse_args(argv=None):
    """
    Parses arguments:
      - If argv is None, reads from the CLI (sys.argv)
      - If argv is an empty list [], uses the defaults
    """
    obj_parser = argparse.ArgumentParser(description='Tissue detection for IF images')

    obj_parser.add_argument('--config', type=str, default=None, help='Path to TOML config file (overrides all other args)')
    obj_parser.add_argument('--image_fullpath', type=str, default=None, help='Input image fullpath (OME-TIFF)')
    obj_parser.add_argument('--channels', type=sp.csv_arg, default=None, help='Comma-separated channel names to use (default: all non-AF channels)')
    obj_parser.add_argument('--exclude_channels', type=sp.csv_arg, default=['AF'], help='Comma-separated channels to exclude from max-projection (default: AF)')
    obj_parser.add_argument('--downsample', type=int, default=4, help='Downsampling factor (default: 4)')
    obj_parser.add_argument('--sigma', type=float, default=1.0, help='Gaussian blur sigma (default: 1.0)')
    obj_parser.add_argument('--threshold', type=str, default='otsu', help='Threshold value or "otsu" (default: otsu)')
    obj_parser.add_argument('--min_obj_size', type=int, default=2500, help='Min object size in pixels at downsampled resolution (default: 2500)')
    obj_parser.add_argument('--max_hole_size', type=int, default=2500, help='Max hole size to fill in pixels at downsampled resolution (default: 2500)')
    obj_parser.add_argument('--output_dir', type=str, default='./', help='Output directory')
    obj_parser.add_argument('--overwrite', action='store_true', help='Overwrite existing output')

    return obj_parser.parse_args(argv)


def fn_main(args):
    """Run tissue detection from parsed CLI arguments or TOML config."""
    print(f'sp-pipeline v{sp.__version__} — tissue detection')
    # If a TOML config is provided, load it and override CLI args
    if (args.config is not None):
        dct_config = fn_load_config_toml(os.path.abspath(args.config))
        str_image_fullpath = dct_config['image_fullpath']
        lst_channels = dct_config.get('channels', None)
        lst_exclude_channels = dct_config.get('exclude_channels', ['AF'])
        int_downsample = dct_config.get('downsample', 4)
        flt_sigma = dct_config.get('sigma', 1.0)
        str_threshold = dct_config.get('threshold', 'otsu')
        int_min_obj_size = dct_config.get('min_obj_size', 2500)
        int_max_hole_size = dct_config.get('max_hole_size', 2500)
        str_output_dir = dct_config.get('output_dir', './')
        bln_overwrite = dct_config.get('overwrite', True)
    else:
        str_image_fullpath = args.image_fullpath
        lst_channels = args.channels
        lst_exclude_channels = args.exclude_channels
        int_downsample = args.downsample
        flt_sigma = args.sigma
        str_threshold = args.threshold
        int_min_obj_size = args.min_obj_size
        int_max_hole_size = args.max_hole_size
        str_output_dir = args.output_dir
        bln_overwrite = args.overwrite

    # Resolve paths to absolute
    str_image_fullpath = os.path.abspath(str_image_fullpath)
    str_output_dir = os.path.abspath(str_output_dir)

    # Parse threshold — could be 'otsu' or a numeric value
    if (str_threshold != 'otsu'):
        try:
            str_threshold = float(str_threshold)
        except ValueError:
            raise ValueError(f'Invalid threshold value: {str_threshold}. Use "otsu" or a number.')

    str_output_dir = os.path.join(str_output_dir, 'tissue_mask')
    os.makedirs(str_output_dir, exist_ok=True)

    str_base_name = sp.fn_get_base_image_file_name(str_image_fullpath)
    str_mask_fullpath = os.path.join(str_output_dir, str_base_name + '_tissue_mask.ome.tiff')
    str_area_fullpath = os.path.join(str_output_dir, str_base_name + '_tissue_area.parquet')

    print(f'Tissue detection: {str_base_name}')
    print(f'  Image: {str_image_fullpath}')

    arr_mask, flt_tissue_area_um2 = fn_detect_tissue(
        arg_image_fullpath=str_image_fullpath,
        arg_output_fullpath=str_mask_fullpath,
        arg_channels=lst_channels,
        arg_exclude_channels=lst_exclude_channels,
        arg_downsample=int_downsample,
        arg_sigma=flt_sigma,
        arg_threshold=str_threshold,
        arg_min_obj_size=int_min_obj_size,
        arg_max_hole_size=int_max_hole_size,
        arg_overwrite=bln_overwrite,
    )

    if (arr_mask is not None):
        # Save tissue area as parquet
        int_tissue_pixels = int(np.sum(arr_mask > 0))
        dct_area = {
            'image_name': [str_base_name],
            'tissue_area_px': [int_tissue_pixels],
            'tissue_area_um2': [flt_tissue_area_um2 if (flt_tissue_area_um2 is not None) else np.nan],
        }
        df_area = pd.DataFrame(dct_area)
        df_area.to_parquet(str_area_fullpath)
        print(f'  Tissue area saved: {str_area_fullpath}')


# ============================================================================
#  Entry point
# ============================================================================

# Example CLI:
#   python 01_tissue_detection.py --image_fullpath './image.ome.tiff'
#   python 01_tissue_detection.py --image_fullpath './image.ome.tiff' --channels DAPI
#   python 01_tissue_detection.py --image_fullpath './image.ome.tiff' --threshold 100
#
# Example TOML:
#   python 01_tissue_detection.py --config config.toml

if __name__ == '__main__':
    args = fn_parse_args()
    fn_main(args)