extract_embedding.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Fri Jan 24 11:19:56 2025

@author: Dr Binghao Chai
@institute: University College London (UCL)

Extract tile-level embeddings from WSIs with Google Path Foundation (TensorFlow).

The script expects:
1. A directory of whole-slide images (e.g. ``.svs``, ``.ndpi``, ``.tif``, ``.mrxs``).
2. A directory of per-slide tile coordinate files (``.h5``) generated by CLAM.

For each matched slide/coordinate pair, the script writes one CSV containing:
``x_coord``, ``y_coord``, and embedding columns (``emb_1 ... emb_n``).

CLI Arguments
-------------
--wsi_path : str
    Directory containing source whole-slide images.
--tile_path : str
    Directory containing CLAM tile coordinate ``.h5`` files.
--model_path : str
    Path to the Path Foundation exported model.
--output_path : str
    Output directory for per-slide embedding CSV files.
"""

# package import
import time
import argparse
from tqdm import tqdm
import os

#import numpy as np
import pandas as pd
from PIL import Image
#import h5py

from openslide import OpenSlide

import tensorflow as tf
from tensorflow.keras.layers import TFSMLayer

from utils.helper_class_tf import WholeSlideBagTF
from utils.helper_functions_tf import resize_transforms

# Reduce XLA register pressure
os.environ["XLA_FLAGS"] = "--xla_gpu_autotune_level=0"

# Limit GPU memory growth to prevent out-of-memory errors
os.environ["TF_FORCE_GPU_ALLOW_GROWTH"] = "true"

# -----------------------------------------------------------------------------
# Slide Processing Function
# -----------------------------------------------------------------------------

def process_slide(tile_path, wsi_path, model, batch_size=32):
    """
    Run patch inference for one slide and collect embedding rows.

    Parameters
    ----------
    tile_path : str
        Path to the CLAM coordinate H5 file for a single slide.
    wsi_path : str
        Path to the source whole-slide image.
    model : tensorflow.keras.layers.TFSMLayer
        Loaded Path Foundation model layer.
    batch_size : int, default=32
        Number of patches processed per inference batch.

    Returns
    -------
    tuple[list[list[float]], int]
        A list of output rows ``[x, y, emb_1, ... emb_n]`` and embedding dimension.
    """
    
    wsi = OpenSlide(wsi_path)
    
    # load tiles
    dataset = WholeSlideBagTF(file_path=tile_path, wsi=wsi, img_transforms=resize_transforms)
    tf_dataset = dataset.get_tf_dataset().batch(batch_size).prefetch(tf.data.AUTOTUNE)
    
    # # Iterate over the dataset - debug print
    # # Take one batch from the tf_dataset
    # for img_batch, coord_batch in tf_dataset.take(1):  # Take one batch
    #     # Print the shape and data type of the batch
    #     print(f"Image batch shape: {img_batch.shape}")  # Should be (batch_size, 224, 224, 3)
    #     print(f"Image batch dtype: {img_batch.dtype}")  # Should be float32
    #     print(f"Coordinate batch shape: {coord_batch.shape}")  # Should be (batch_size, 2)
    #     print(f"Coordinate batch dtype: {coord_batch.dtype}")  # Should be int32
    
    #     # Check the pixel value range
    #     print(f"Image pixel value range: {tf.reduce_min(img_batch).numpy()} to {tf.reduce_max(img_batch).numpy()}")
    
    #     # Optionally inspect a specific image and its coordinates
    #     print("\nFirst image and its coordinates:")
    #     print(img_batch[0].numpy())  # Convert to NumPy for inspection
    #     print(coord_batch[0].numpy())  # Convert to NumPy for inspection
    
    # Initialize an empty list to collect data for the DataFrame    
    outputs = []
    
    for img_batch, coord_batch in tqdm(tf_dataset, desc="Processing Batches"):
        result = model(img_batch)
        batch_output = result['output_0'].numpy()
        coords_np = coord_batch.numpy()
        
        for i in range(batch_output.shape[0]):
            outputs.append([coords_np[i][0], coords_np[i][1]] + batch_output[i].tolist())
            
    nr_features = batch_output.shape[1]
    
    return outputs, nr_features

# -----------------------------------------------------------------------------
# Main Entry Point
# -----------------------------------------------------------------------------
if __name__ == "__main__":
    # define argument parser
    parser = argparse.ArgumentParser(
        description="Extract tile embeddings from WSIs using Google Path Foundation."
    )
    parser.add_argument('--wsi_path', type=str, required=True,
    					help='Directory containing source whole-slide images.')
    parser.add_argument('--tile_path', type=str, required=True, 
    					help='Directory containing CLAM tile coordinate .h5 files.')
    parser.add_argument('--model_path', type=str, required=True, 
    					help='Path to the Path Foundation exported model.')
    parser.add_argument('--output_path', type=str, required=True,
                        help='Output directory for per-slide embedding CSV files.')
    args = parser.parse_args()

    # Print GPU availability
    print("Num GPUs Available:", len(tf.config.list_physical_devices('GPU')))
    
    since = time.time()
    
    # Path to the pre-trained model (TFSMLayer)
    model_path = args.model_path
    path_foundation = TFSMLayer(model_path, call_endpoint="serving_default")
    
    # Ensure output directory exists
    os.makedirs(args.output_path, exist_ok=True)
    
    # List existing embedding files (without extensions)
    embedding_files = set(f.replace(".csv", "") for f in os.listdir(args.output_path) if f.endswith(".csv"))
    
    # List all WSI and tile files
    wsi_files = sorted([f for f in os.listdir(args.wsi_path) if f.endswith((".svs", ".ndpi", ".tif", ".mrxs"))])
    tile_files = sorted([f for f in os.listdir(args.tile_path) if f.endswith(".h5")])
    
    # Filter out WSIs and tiles that already have corresponding embeddings
    filtered_wsi_files = []
    filtered_tile_files = []
    
    for wsi_file, tile_file in zip(wsi_files, tile_files):
        base_name = os.path.splitext(wsi_file)[0]  # Get filename without extension
        if base_name not in embedding_files:  # If embedding doesn't exist, add to processing list
            filtered_wsi_files.append(wsi_file)
            filtered_tile_files.append(tile_file)
    
    print(f"Skipping {len(wsi_files) - len(filtered_wsi_files)} slides that already have embeddings.")
    print(f"Processing {len(filtered_wsi_files)} slides.")

    for wsi_file, tile_file in tqdm(zip(filtered_wsi_files, filtered_tile_files), desc="Processing Slides", total=len(filtered_wsi_files)):

        # Ensure the filenames match (excluding extensions)
        if os.path.splitext(wsi_file)[0] != os.path.splitext(tile_file)[0]:
            print(f"Skipping {wsi_file} and {tile_file}: filenames do not match.")
            continue
    
        # Construct full paths
        wsi_path = os.path.join(args.wsi_path, wsi_file)
        tile_path = os.path.join(args.tile_path, tile_file)
    
        print(f"Processing slide: {wsi_file}")
    
        # Process the slide
        outputs, nr_features = process_slide(
            tile_path=tile_path,
            wsi_path=wsi_path,
            model=path_foundation,
            batch_size=32,
        )
    
        # Create a DataFrame with the required columns
        columns = ['x_coord', 'y_coord'] + [f'emb_{i+1}' for i in range(nr_features)]
        df = pd.DataFrame(outputs, columns=columns)
    
        # Save the DataFrame to a CSV file
        slide_name = os.path.splitext(wsi_file)[0]
        output_file = os.path.join(args.output_path, f"{slide_name}.csv")
        df.to_csv(output_file, index=False)
        print(f"Features saved to {output_file}")
    
    print("All slides processed!")
        
    # outputs, nr_features = process_slide(
    #     tile_path="/home/digitalpathology/workspace/path_foundation_stain_variation_refactoring/tiles/01_external/patches/RNOH_1_3_S00140422_131427.h5",
    #     wsi_path="/home/digitalpathology/workspace/path_foundation_stain_variation_refactoring/source_wsi/01_external/RNOH_1_3_S00140422_131427.svs",
    #     model=path_foundation,
    #     batch_size=32,
    #     )
    
    # # Create a DataFrame with the required columns
    # columns = ['x_coord', 'y_coord'] + [f'emb_{i+1}' for i in range(nr_features)]
    # df = pd.DataFrame(outputs, columns=columns)
    
    # # Save the DataFrame to a single CSV file
    # df.to_csv("output_features_with_coords.csv", index=False)
    
    # print("Feature extraction and saving completed!")
    
    # Print the total runtime
    time_elapsed = time.time() - since
    print("Task complete in {:.0f}m {:.0f}s".format(time_elapsed // 60, time_elapsed % 60))