extract_features.py

import argparse
import os
import time
import numpy as np

import openslide
import cv2
from PIL import Image, ImageDraw
from shapely.affinity import scale
from shapely.geometry import Polygon, MultiPolygon
from shapely.ops import unary_union
from collections import defaultdict

import nmslib
from sklearn.metrics.pairwise import euclidean_distances, cosine_similarity

# Optional if stain deconvolution is used.
import histomicstk as htk #pip install histomicstk --find-links https://girder.github.io/large_image_wheels

import torch
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from esvit.utils import bool_flag

# You can use your own encoder to extract features. Here's examples including EsVIT the one used in the publication. 
from encoders import load_encoder_esVIT, load_encoder_resnet

def get_args_parser():
    parser = argparse.ArgumentParser('Preprocessing script esvit', add_help=False)
    parser.add_argument(
        "--input_slide",
        type=str,
        help="Path to input WSI file",
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        help="Directory to save output data",
    )
    parser.add_argument(
        "--checkpoint",
        type=str,
        help="Feature extractor weights checkpoint",
    )
    parser.add_argument(
        "--batch_size",
        type=int,
        default=512,
    )
    parser.add_argument(
        "--tile_size",
        help="Desired tile size in microns (should be the same value as used in feature extraction model).",
        type=int,
        required=True,
    )
    parser.add_argument(
        "--out_size",
        help="Resize the square tile to this output size (in pixels).",
        type=int,
        default=224,
    )
    parser.add_argument(
        "--method",
        help="Segmentation method, otsu or stain deconv",
        type=str,
        default='otsu',
    )
    parser.add_argument(
        "--dist_threshold",
        type=int,
        default=4,
        help="L2 norm distance when spatially merging pacthes.",
    )
    parser.add_argument(
        "--corr_threshold",
        type=float,
        default=0.6,
        help="Cosine similarity distance when semantically merging pacthes.",
    )
    parser.add_argument(
        "--workers",
        help="The number of workers to use for the data loader. Only relevant when using a GPU.",
        type=int,
        default=4,
    )
    parser.add_argument(
        '--cfg',
        help='experiment configure file name. See EsVIT repo.',
        type=str
    )
    parser.add_argument(
        '--arch', default='deit_small', type=str,
        choices=['cvt_tiny', 'swin_tiny','swin_small', 'swin_base', 'swin_large', 'swin', 'vil', 'vil_1281', 'vil_2262', 'deit_tiny', 'deit_small', 'vit_base'],
        help="""Name of architecture to train. For quick experiments with ViTs, we recommend using deit_tiny or deit_small. See EsVIT repo."""
    )
    parser.add_argument(
        '--n_last_blocks', 
        default=4, 
        type=int, 
        help="""Concatenate [CLS] tokens for the `n` last blocks. We use `n=4` when evaluating DeiT-Small and `n=1` with ViT-Base. See EsVIT repo."""
    )
    parser.add_argument(
        '--avgpool_patchtokens', 
        default=False, 
        type=bool_flag,
        help="""Whether ot not to concatenate the global average pooled features to the [CLS] token.
        We typically set this to False for DeiT-Small and to True with ViT-Base. See EsVIT repo."""
    )
    parser.add_argument(
        '--patch_size', 
        default=8, 
        type=int, 
        help='Patch resolution of the model. See EsVIT repo.'
    )
    parser.add_argument(
        'opts',
        help="Modify config options using the command-line. See EsVIT repo.",
        default=None,
        nargs=argparse.REMAINDER
    )
    parser.add_argument(
        "--rank", 
        default=0, 
        type=int, 
        help="Please ignore and do not set this argument.")

    return parser

def segment_tissue(img):
    img_hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    mthresh = 7
    img_med = cv2.medianBlur(img_hsv[:, :, 1], mthresh)
    _, img_prepped = cv2.threshold(img_med, 0, 255, cv2.THRESH_OTSU + cv2.THRESH_BINARY)

    close = 4
    kernel = np.ones((close, close), np.uint8)
    img_prepped = cv2.morphologyEx(img_prepped, cv2.MORPH_CLOSE, kernel)

    # Find and filter contours
    contours, hierarchy = cv2.findContours(
        img_prepped, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE
    )
    return contours, hierarchy

def segment_tissue_deconv_stain(img):
    """
    Method 2: Tissue segmentation using stain deconvolution. Alternative to Otsu thresholding. 
    """
    image = img.copy()

    image[image[...,-1]==0] = [255,255,255,0]

    image = Image.fromarray(image)
    image = np.asarray(image.convert('RGB'))
    
    I_0 = 255
    
    # Create stain to color map
    stain_color_map = htk.preprocessing.color_deconvolution.stain_color_map

    # Specify stains of input image
    stains = ['hematoxylin',  # nuclei stain
              'eosin']        # cytoplasm stain
    
    w_est = htk.preprocessing.color_deconvolution.rgb_separate_stains_macenko_pca(image, I_0)
    deconv_result = htk.preprocessing.color_deconvolution.color_deconvolution(image, w_est, I_0)
    
    final_mask = np.zeros(image.shape[0:2], np.uint8)

    for i in 0, 1: 
        channel = htk.preprocessing.color_deconvolution.find_stain_index(
            stain_color_map[stains[i]], w_est)

        img_for_thresholding = 255 - deconv_result.Stains[:, :, channel]
        _, img_prepped = cv2.threshold(
            img_for_thresholding, 0, 255, cv2.THRESH_OTSU+cv2.THRESH_BINARY)

        final_mask = cv2.bitwise_or(final_mask, img_prepped)
        
    for i in range(5):
        close = 3
        kernel = np.ones((close, close), np.uint8)
        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_OPEN, kernel)
        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_CLOSE, kernel)
    
    return final_mask

def mask_to_polygons(mask, min_area, min_area_holes=10., epsilon=10.):
    """Convert a mask ndarray (binarized image) to Multipolygons"""
    # first, find contours with cv2: it's much faster than shapely
    contours, hierarchy = cv2.findContours(mask,
                                  cv2.RETR_CCOMP,
                                  cv2.CHAIN_APPROX_NONE)
    if not contours:
        return MultiPolygon()

    cnt_children = defaultdict(list)
    child_contours = set()
    assert hierarchy.shape[0] == 1
    
    # http://docs.opencv.org/3.1.0/d9/d8b/tutorial_py_contours_hierarchy.html
    for idx, (_, _, _, parent_idx) in enumerate(hierarchy[0]):
        if parent_idx != -1:
            child_contours.add(idx)
            cnt_children[parent_idx].append(contours[idx])
            
    # create actual polygons filtering by area (removes artifacts)
    all_polygons = []

    for idx, cnt in enumerate(contours):

        if idx not in child_contours and cv2.contourArea(cnt) >= min_area:
            assert cnt.shape[1] == 1
            poly = Polygon(
                shell=cnt[:, 0, :],
                holes=[c[:, 0, :] for c in cnt_children.get(idx, [])
                       if cv2.contourArea(c) >= min_area_holes])
            
            if not poly.is_valid:
                # This is likely becausee the polygon is self-touching or self-crossing.
                # Try and 'correct' the polygon using the zero-length buffer() trick.
                # See https://shapely.readthedocs.io/en/stable/manual.html#object.buffer
                poly = poly.buffer(0)
    
            all_polygons.append(poly)

    if len(all_polygons) == 0:
        raise Exception("Raw tissue mask consists of 0 polygons")

    # if this raises an issue - try instead unary_union(all_polygons)        
    all_polygons = MultiPolygon(all_polygons)

    return all_polygons

def detect_foreground(contours, hierarchy):
    hierarchy = np.squeeze(hierarchy, axis=(0,))[:, 2:]

    # find foreground contours (parent == -1)
    hierarchy_1 = np.flatnonzero(hierarchy[:, 1] == -1)
    foreground_contours = [contours[cont_idx] for cont_idx in hierarchy_1]

    all_holes = []
    for cont_idx in hierarchy_1:
        all_holes.append(np.flatnonzero(hierarchy[:, 1] == cont_idx))

    hole_contours = []
    for hole_ids in all_holes:
        holes = [contours[idx] for idx in hole_ids]
        hole_contours.append(holes)

    return foreground_contours, hole_contours

def construct_polygon(foreground_contours, hole_contours, min_area):
    polys = []
    for foreground, holes in zip(foreground_contours, hole_contours):
        # We remove all contours that consist of fewer than 3 points, as these won't work with the Polygon constructor.
        if len(foreground) < 3:
            continue

        # remove redundant dimensions from the contour and convert to Shapely Polygon
        poly = Polygon(np.squeeze(foreground))

        # discard all polygons that are considered too small
        if poly.area < min_area:
            continue

        if not poly.is_valid:
            # This is likely becausee the polygon is self-touching or self-crossing.
            # Try and 'correct' the polygon using the zero-length buffer() trick.
            # See https://shapely.readthedocs.io/en/stable/manual.html#object.buffer
            poly = poly.buffer(0)

        # Punch the holes in the polygon
        for hole_contour in holes:
            if len(hole_contour) < 3:
                continue

            hole = Polygon(np.squeeze(hole_contour))

            if not hole.is_valid:
                continue

            # ignore all very small holes
            if hole.area < min_area:
                continue

            poly = poly.difference(hole)

        polys.append(poly)

    if len(polys) == 0:
        raise Exception("Raw tissue mask consists of 0 polygons")

    # If we have multiple polygons, we merge any overlap between them using unary_union().
    # This will result in a Polygon or MultiPolygon with most tissue masks.
    return unary_union(polys)

def generate_tiles(tile_width_pix, tile_height_pix, img_width, img_height, offsets=[(0, 0)]):
    # Generate tiles covering the entire image.
    # Provide an offset (x,y) to create a stride-like overlap effect.
    # Add an additional tile size to the range stop to prevent tiles being cut off at the edges.
    range_stop_width = int(np.ceil(img_width + tile_width_pix))
    range_stop_height = int(np.ceil(img_height + tile_height_pix))

    rects = []
    for xmin, ymin in offsets:
        cols = range(int(np.floor(xmin)), range_stop_width, tile_width_pix)
        rows = range(int(np.floor(ymin)), range_stop_height, tile_height_pix)
        for x in cols:
            for y in rows:
                rect = Polygon(
                    [
                        (x, y),
                        (x + tile_width_pix, y),
                        (x + tile_width_pix, y - tile_height_pix),
                        (x, y - tile_height_pix),
                    ]
                )
                rects.append(rect)
    return rects

def make_tile_QC_fig(tiles, slide, level, line_width_pix=1, extra_tiles=None):
    # Render the tiles on an image derived from the specified zoom level
    img = slide.read_region((0, 0), level, slide.level_dimensions[level])
    downsample = 1 / slide.level_downsamples[level]

    draw = ImageDraw.Draw(img, "RGBA")
    for tile in tiles:
        bbox = tuple(np.array(tile.bounds) * downsample)
        draw.rectangle(bbox, outline="lightgreen", width=line_width_pix)

    # allow to display other tiles, such as excluded or sampled
    if extra_tiles:
        for tile in extra_tiles:
            bbox = tuple(np.array(tile.bounds) * downsample)
            draw.rectangle(bbox, outline="blue", width=line_width_pix + 1)

    return img

def create_tissue_mask(wsi, seg_level, method='otsu'):
    # Determine the best level to determine the segmentation on
    level_dims = wsi.level_dimensions[seg_level]

    img = np.array(wsi.read_region((0, 0), seg_level, level_dims))

    # Get the total surface area of the slide level that was used
    level_area = level_dims[0] * level_dims[1]

    # Minimum surface area of tissue polygons (in pixels)
    # Note that this value should be sensible in the context of the chosen tile size
    min_area = level_area / 500

    if method=='stain_deconv':
        tissue_mask = segment_tissue_deconv_stain(img)
        tissue_mask = mask_to_polygons(tissue_mask, min_area)
    else:
        contours, hierarchy = segment_tissue(img)
        foreground_contours, hole_contours = detect_foreground(contours, hierarchy)
        tissue_mask = construct_polygon(foreground_contours, hole_contours, min_area)

    # Scale the tissue mask polygon to be in the coordinate space of the slide's level 0
    scale_factor = wsi.level_downsamples[seg_level]
    tissue_mask_scaled = scale(
        tissue_mask, xfact=scale_factor, yfact=scale_factor, zfact=1.0, origin=(0, 0)
    )

    return tissue_mask_scaled

def create_tissue_tiles(wsi, tissue_mask_scaled, tile_size_microns, offsets_micron=None):

    print(f"tile size is {tile_size_microns} um")

    # Compute the tile size in pixels from the desired tile size in microns and the image resolution
    assert (
        openslide.PROPERTY_NAME_MPP_X in wsi.properties
    ), "microns per pixel along X-dimension not available"
    assert (
        openslide.PROPERTY_NAME_MPP_Y in wsi.properties
    ), "microns per pixel along Y-dimension not available"

    mpp_x = float(wsi.properties[openslide.PROPERTY_NAME_MPP_X])
    mpp_y = float(wsi.properties[openslide.PROPERTY_NAME_MPP_Y])

    # For larger tiles in micron, NKI scanner outputs mppx slight different than mppy.
    # Force tiles to be squared.
    mpp_scale_factor = min(mpp_x, mpp_y)
    if mpp_x != mpp_y:
        print(
            f"mpp_x of {mpp_x} and mpp_y of {mpp_y} are not the same. Using smallest value: {mpp_scale_factor}"
        )

    tile_size_pix = round(tile_size_microns / mpp_scale_factor)

    # Use the tissue mask bounds as base offsets (+ a margin of a few tiles) to avoid wasting CPU power creating tiles that are never going
    # to be inside the tissue mask.
    tissue_margin_pix = tile_size_pix * 2
    minx, miny, maxx, maxy = tissue_mask_scaled.bounds
    min_offset_x = minx - tissue_margin_pix
    min_offset_y = miny - tissue_margin_pix
    offsets = [(min_offset_x, min_offset_y)]

    if offsets_micron is not None:
        assert (
            len(offsets_micron) > 0
        ), "offsets_micron needs to contain at least one value"
        # Compute the offsets in micron scale
        offset_pix = [round(o / mpp_scale_factor) for o in offsets_micron]
        offsets = [(o + min_offset_x, o + min_offset_y) for o in offset_pix]

    # Generate tiles covering the entire WSI
    all_tiles = generate_tiles(
        tile_size_pix,
        tile_size_pix,
        maxx + tissue_margin_pix,
        maxy + tissue_margin_pix,
        offsets=offsets,
    )

    # Retain only the tiles that sit within the tissue mask polygon
    filtered_tiles = [rect for rect in all_tiles if tissue_mask_scaled.intersects(rect)]

    return filtered_tiles

def tile_is_not_empty(tile, threshold_white=20):
    histogram = tile.histogram()

    # Take the median of each RGB channel. Alpha channel is not of interest.
    # If roughly each chanel median is below a threshold, i.e close to 0 till color value around 250 (white reference) then tile mostly white.
    whiteness_check = [0, 0, 0]
    for channel_id in (0, 1, 2):
        whiteness_check[channel_id] = np.median(
            histogram[256 * channel_id : 256 * (channel_id + 1)][100:200]
        )

    if all(c <= threshold_white for c in whiteness_check):
        # exclude tile
        return False

    # keep tile
    return True

def crop_rect_from_slide(slide, rect):
    minx, miny, maxx, maxy = rect.bounds
    # Note that the y-axis is flipped in the slide: the top of the shapely polygon is y = ymax,
    # but in the slide it is y = 0. Hence: miny instead of maxy.
    top_left_coords = (int(minx), int(miny))
    return slide.read_region(top_left_coords, 0, (int(maxx - minx), int(maxy - miny)))

class BagOfTiles(Dataset):
    def __init__(self, wsi, tiles, resize_to=224):
        self.wsi = wsi
        self.tiles = tiles

        self.roi_transforms = transforms.Compose(
            [
                # As we can't be sure that the input tile dimensions are all consistent, we resize
                # them to a commonly used size before feeding them to the model.
                # Note: assumes a square image.
                transforms.Resize(resize_to),
                # Turn the PIL image into a (C x H x W) float tensor in the range [0.0, 1.0]
                transforms.ToTensor(),
            ]
        )

    def __len__(self):
        return len(self.tiles)

    def __getitem__(self, idx):
        tile = self.tiles[idx]
        img = crop_rect_from_slide(self.wsi, tile)

        # RGB filtering - calling here speeds up computation since it requires crop_rect_from_slide function.
        #is_tile_kept = tile_is_not_empty(img, threshold_white=20)
        is_tile_kept = True

        # Ensure the img is RGB, as expected by the pretrained model.
        # See https://pytorch.org/docs/stable/torchvision/models.html
        img = img.convert("RGB")

        # Ensure we have a square tile in our hands.
        # We can't handle non-squares currently, as this would requiring changes to
        # the aspect ratio when resizing.
        width, height = img.size
        assert width == height, "input image is not a square"

        img = self.roi_transforms(img).unsqueeze(0)
        coord = tile.bounds
        return img, coord, is_tile_kept

def collate_features(batch):
    # Item 2 is the boolean value from tile filtering.
    img = torch.cat([item[0] for item in batch if item[2]], dim=0)
    coords = np.vstack([item[1] for item in batch if item[2]])
    return [img, coords]

def mergedpatch_gen(features, coords, dist_threshold=4, corr_threshold = 0.6):

    # Get patch distance in pixels with rendered segmentation level. Note that each patch is squared and therefore same distance.
    patch_dist = abs(coords[0,2] - coords[0,0]) 
    print(patch_dist)
    
    # Compute feature similarity (cosine) and nearby pacthes (L2 norm - only need the top left x,y coordinates)
    cosine_matrix = cosine_similarity(features, features)
    coordinate_matrix = euclidean_distances(coords[:,:2], coords[:,:2])

    # NOTE: random selection for the first patch for patch merging might be less biased towards tissue orientation and size. 
    indices_avail = np.arange(features.shape[0])
    np.random.seed(0)  
    np.random.shuffle(indices_avail)

    # Merging together nearby patches and similar within pre-defined threshold. 
    mergedfeatures = []
    indices_used = []
    for ref in indices_avail:

        # This has been merged already
        if ref not in indices_used:

            # Making sure they won't be selected once more
            if indices_used:
                coordinate_matrix[ref,indices_used] = [np.Inf]*len(indices_used)
                cosine_matrix[ref,indices_used] = [0.0]*len(indices_used)
            
            indices_dist = np.where(coordinate_matrix[ref] < patch_dist*dist_threshold, 1 , 0)
            indices_corr = np.where(cosine_matrix[ref] > corr_threshold, 1 , 0)
            final_indices = indices_dist * indices_corr

            # which includes already the ref patch
            indices_used.extend(list(np.where(final_indices == 1)[0]))
            mergedfeatures.append(tuple((features[final_indices==1,:], coords[final_indices==1,:])))
        else:
            continue
        
    assert len(indices_used)==features.shape[0], f'Probably issue in contruscting merged features for graph {len(indices_used)}!={features.shape[0]}'

    return mergedfeatures

class HNSW:
    def __init__(self, space):
        self.space = space

    def fit(self, X):
        # See https://nmslib.github.io/nmslib/quickstart.html
        index = nmslib.init(space=self.space, method='hnsw')
        index.addDataPointBatch(X)
        index.createIndex()
        self.index_ = index
        return self

    def query(self, vector, topn):
        indices, dist = self.index_.knnQuery(vector, k=topn)
        return indices, dist

@torch.no_grad()
def extract_features(model, device, wsi, filtered_tiles, workers, out_size, batch_size, n_last_blocks, avgpool_patchtokens, depths):
    # Use multiple workers if running on the GPU, otherwise we'll need all workers for evaluating the model.
    kwargs = (
        {"num_workers": workers, "pin_memory": True} if device.type == "cuda" else {}
    )
    loader = DataLoader(
        dataset=BagOfTiles(wsi, filtered_tiles, resize_to=out_size),
        batch_size=batch_size,
        collate_fn=collate_features,
        **kwargs,
    )
    features_ = []
    coords_ = []
    for batch, coords in loader:
        batch = batch.to(device, non_blocking=True)
        # NOTE: Example using EsVIT. You may want to call your own feature extractor otherwise. 
        features = model.forward_return_n_last_blocks(batch, n_last_blocks, avgpool_patchtokens, depths).cpu().numpy()
        features_.extend(features)
        coords_.extend(coords)
    return np.asarray(features_), np.asarray(coords_)

def extract_save_features(args):
    # Derive the slide ID from its name.
    slide_id, _ = os.path.splitext(os.path.basename(args.input_slide))
    wip_file_path = os.path.join(args.output_dir, slide_id + "_wip.h5")
    output_file_path = os.path.join(args.output_dir, slide_id + "_features.h5")

    os.makedirs(args.output_dir, exist_ok=True)

    # Check if the _features output file already exist. If so, we terminate to avoid
    # overwriting it by accident. This also simplifies resuming bulk batch jobs.
    if os.path.exists(output_file_path):
        raise Exception(f"{output_file_path} already exists")

    # Open the slide for reading.
    wsi = openslide.open_slide(args.input_slide)

    # Decide on which slide level we want to base the segmentation.
    seg_level = wsi.get_best_level_for_downsample(64)

    # Run the segmentation and  tiling procedure.
    start_time = time.time()
    tissue_mask_scaled = create_tissue_mask(wsi, seg_level, method=args.method)
    filtered_tiles = create_tissue_tiles(wsi, tissue_mask_scaled, args.tile_size)

    # Build a figure for quality control purposes, to check if the tiles are where we expect them.
    qc_img = make_tile_QC_fig(filtered_tiles, wsi, seg_level, 2)
    qc_img_target_width = 1920
    qc_img = qc_img.resize((qc_img_target_width, int(qc_img.height / (qc_img.width / qc_img_target_width))))
    qc_img_file_path = os.path.join(args.output_dir, f"{slide_id}_features_QC.png")
    qc_img.save(qc_img_file_path)
    print(f"Finished creating {len(filtered_tiles)} tissue tiles in {time.time() - start_time}s")

    # Save QC figure.
    qc_img_file_path = os.path.join(
        args.output_dir, f"{slide_id}_N{len(mergedpatches)}mergedpatches_distThreshold{args.dist_threshold}_corrThreshold{args.corr_threshold}.png"
    )

    # Extract the rectangles, and compute the feature vectors. Example using EsVIT. 
    device = torch.device("cuda") 
    model, _, depths = load_encoder_esVIT(args, device)
    
    features, coords = extract_features(
        model,
        device,
        wsi,
        filtered_tiles,
        args.workers,
        args.out_size,
        args.batch_size,
        n_last_blocks = args.n_last_blocks, 
        avgpool_patchtokens = args.avgpool_patchtokens,
        depths = depths,
    )
    
    print(f'Number of features N={len(features)}')
    # Merging nearby patches with similar semantic. 
    mergedpatches = mergedpatch_gen(features, coords, dist_threshold=args.dist_threshold, corr_threshold=args.corr_threshold)
    print(f'Merging step => N={len(mergedpatches)}')

    # Saving features.
    torch.save(mergedpatches, wip_file_path)

    # Rename the file containing the patches to ensure we can easily
    # distinguish incomplete bags of patches (due to e.g. errors) from complete ones in case a job fails.
    os.rename(wip_file_path, output_file_path)

    print('Done.')

if __name__ == '__main__':
    parser = argparse.ArgumentParser('Preprocessing script esvit', parents=[get_args_parser()])
    args = parser.parse_args()

    assert os.path.isfile(args.checkpoint), f'{args.checkpoint} does not exist'
    assert torch.cuda.is_available(), 'Need cuda for this job'
    assert os.path.isfile(args.input_slide), f'{args.input_slide} does not exist'

    extract_save_features(args)