create_patches_fp.py

# internal imports
from wsi_core.WholeSlideImage import WholeSlideImage
from wsi_core.wsi_utils import StitchCoords
from wsi_core.batch_process_utils import initialize_df
# other imports
import os
import numpy as np
import time
import argparse
import pandas as pd
import fsspec
from pathlib import Path

VALID_SLIDE_EXTENSIONS = [".svs", ".scn", ".tif"]

def stitching(file_path, wsi_object, downscale = 64):
    start = time.time()
    heatmap = StitchCoords(file_path, wsi_object, 
                        downscale=downscale, bg_color=(0,0,0), 
                        alpha=-1, draw_grid=False)
    total_time = time.time() - start
    
    return heatmap, total_time

def segment(WSI_object, seg_params = None, filter_params = None, mask_file = None):
    ### Start Seg Timer
    start_time = time.time()
    # Use segmentation file
    if mask_file is not None:
        WSI_object.initSegmentation(mask_file)
    # Segment   
    else:
        WSI_object.segmentTissue(**seg_params, filter_params=filter_params)

    ### Stop Seg Timers
    seg_time_elapsed = time.time() - start_time   
    return WSI_object, seg_time_elapsed

def patching(WSI_object, h5_file_path, **kwargs):
    ### Start Patch Timer
    start_time = time.time()

    # Patch
    file_path = WSI_object.process_contours(h5_file_path, **kwargs)


    ### Stop Patch Timer
    patch_time_elapsed = time.time() - start_time
    return file_path, patch_time_elapsed


def seg_and_patch(source, save_dir, patch_save_dir, mask_save_dir, stitch_save_dir, 
                  patch_size = 256, step_size = 256, 
                  seg_params = {'seg_level': -1, 'sthresh': 8, 'mthresh': 7, 'close': 4, 'use_otsu': False,
                  'keep_ids': 'none', 'exclude_ids': 'none'},
                  filter_params = {'a_t':100, 'a_h': 16, 'max_n_holes':8}, 
                  vis_params = {'vis_level': -1, 'line_thickness': 500},
                  patch_params = {'use_padding': True, 'contour_fn': 'four_pt'},
                  patch_level = 0,
                  use_default_params = False, 
                  seg = False, save_mask = True, 
                  stitch= False, 
                  patch = False, auto_skip=True, process_list = None, 
                  s3_storage_key = None,
                  s3_storage_secret = None,
                  s3_endpoint_url = None):
    
    storage_options = {"key" : s3_storage_key, 
                       "secret" : s3_storage_secret, 
                       "client_kwargs" : {'endpoint_url' : s3_endpoint_url}}

    # Create fsspec filesystems for file operations, 
    # simplecache is used for dst_filesystem to allow seeking/writing to h5 files 
    # at the same time since the default s3 filesystem does not support it
    # Url Chaining: https://filesystem-spec.readthedocs.io/en/latest/features.html#url-chaining
    src_filesystem, slide_path = fsspec.core.url_to_fs(source, **storage_options)
    dst_filesystem, dst_path = fsspec.core.url_to_fs("simplecache::" + save_dir, s3=storage_options)
    
    dst_filesystem.makedirs(str(Path(dst_path + patch_save_dir)), exist_ok=True)
    dst_filesystem.makedirs(str(Path(dst_path + mask_save_dir)), exist_ok=True)
    dst_filesystem.makedirs(str(Path(dst_path + stitch_save_dir)), exist_ok=True)

    slides = []  # type: list[str]
    
    if any([slide_path.endswith(ext) for ext in VALID_SLIDE_EXTENSIONS]):
        slides += slide_path
    else:
        for ext in VALID_SLIDE_EXTENSIONS:
            slides += src_filesystem.glob(f"{slide_path}/*{ext}")
    slides = [slide.split('/')[-1] for slide in slides]

    if process_list is None:
        df = initialize_df(slides, seg_params, filter_params, vis_params, patch_params)
    
    else:
        df = pd.read_csv(process_list)
        df = initialize_df(df, seg_params, filter_params, vis_params, patch_params)

    mask = df['process'] == 1
    process_stack = df[mask]

    total = len(process_stack)

    legacy_support = 'a' in df.keys()
    if legacy_support:
        print('detected legacy segmentation csv file, legacy support enabled')
        df = df.assign(**{'a_t': np.full((len(df)), int(filter_params['a_t']), dtype=np.uint32),
        'a_h': np.full((len(df)), int(filter_params['a_h']), dtype=np.uint32),
        'max_n_holes': np.full((len(df)), int(filter_params['max_n_holes']), dtype=np.uint32),
        'line_thickness': np.full((len(df)), int(vis_params['line_thickness']), dtype=np.uint32),
        'contour_fn': np.full((len(df)), patch_params['contour_fn'])})

    seg_times = 0.
    patch_times = 0.
    stitch_times = 0.

    for i in range(total):
        csv_path = str(Path(dst_path) / 'process_list_autogen.csv')
        with dst_filesystem.open(csv_path, 'w') as f:
            df.to_csv(f, index=False)
        idx = process_stack.index[i]
        slide = process_stack.loc[idx, 'slide_id']
        print("\n\nprogress: {:.2f}, {}/{}".format(i/total, i, total))
        print('processing {}'.format(slide))
        
        df.loc[idx, 'process'] = 0
        slide_id, _ = os.path.splitext(slide)

        if auto_skip and dst_filesystem.isfile(str(Path(dst_path + patch_save_dir) / f'{slide_id}.h5')):
            print('{} already exist in destination location, skipped'.format(slide_id))
            df.loc[idx, 'status'] = 'already_exist'
            continue

        # Inialize WSI
        full_path = os.path.join(source, slide)
        WSI_object = WholeSlideImage(full_path, storage_options)

        if use_default_params:
            current_vis_params = vis_params.copy()
            current_filter_params = filter_params.copy()
            current_seg_params = seg_params.copy()
            current_patch_params = patch_params.copy()
            
        else:
            current_vis_params = {}
            current_filter_params = {}
            current_seg_params = {}
            current_patch_params = {}


            for key in vis_params.keys():
                if legacy_support and key == 'vis_level':
                    df.loc[idx, key] = -1
                current_vis_params.update({key: df.loc[idx, key]})

            for key in filter_params.keys():
                if legacy_support and key == 'a_t':
                    old_area = df.loc[idx, 'a']
                    seg_level = df.loc[idx, 'seg_level']
                    scale = WSI_object.level_downsamples[seg_level]
                    adjusted_area = int(old_area * (scale[0] * scale[1]) / (512 * 512))
                    current_filter_params.update({key: adjusted_area})
                    df.loc[idx, key] = adjusted_area
                current_filter_params.update({key: df.loc[idx, key]})

            for key in seg_params.keys():
                if legacy_support and key == 'seg_level':
                    df.loc[idx, key] = -1
                current_seg_params.update({key: df.loc[idx, key]})

            for key in patch_params.keys():
                current_patch_params.update({key: df.loc[idx, key]})

        if current_vis_params['vis_level'] < 0:
            if len(WSI_object.level_dim) == 1:
                current_vis_params['vis_level'] = 0
            
            else:   
                wsi = WSI_object.getOpenSlide()
                best_level = wsi.get_best_level_for_downsample(64)
                current_vis_params['vis_level'] = best_level

        if current_seg_params['seg_level'] < 0:
            if len(WSI_object.level_dim) == 1:
                current_seg_params['seg_level'] = 0
            
            else:
                wsi = WSI_object.getOpenSlide()
                best_level = wsi.get_best_level_for_downsample(64)
                current_seg_params['seg_level'] = best_level

        keep_ids = str(current_seg_params['keep_ids'])
        if keep_ids != 'none' and len(keep_ids) > 0:
            str_ids = current_seg_params['keep_ids']
            current_seg_params['keep_ids'] = np.array(str_ids.split(',')).astype(int)
        else:
            current_seg_params['keep_ids'] = []

        exclude_ids = str(current_seg_params['exclude_ids'])
        if exclude_ids != 'none' and len(exclude_ids) > 0:
            str_ids = current_seg_params['exclude_ids']
            current_seg_params['exclude_ids'] = np.array(str_ids.split(',')).astype(int)
        else:
            current_seg_params['exclude_ids'] = []

        w, h = WSI_object.level_dim[current_seg_params['seg_level']] 
        if w * h > 1e8:
            print('level_dim {} x {} is likely too large for successful segmentation, aborting'.format(w, h))
            df.loc[idx, 'status'] = 'failed_seg'
            continue

        df.loc[idx, 'vis_level'] = current_vis_params['vis_level']
        df.loc[idx, 'seg_level'] = current_seg_params['seg_level']


        seg_time_elapsed = -1
        if seg:
            WSI_object, seg_time_elapsed = segment(WSI_object, current_seg_params, current_filter_params) 

        if save_mask:
            mask = WSI_object.visWSI(**current_vis_params)
            mask_path = str(Path(dst_path + mask_save_dir) / f'{slide_id}.jpg')
            with dst_filesystem.open(mask_path, 'wb') as f:
                mask.save(f, "JPEG")

        patch_time_elapsed = -1 # Default time
        if patch:
            save_path = str(Path(dst_path + patch_save_dir) / f'{slide_id}.h5')
            current_patch_params.update({'patch_level': patch_level, 'patch_size': patch_size, 'step_size': step_size, 
                                        'save_path': save_path})
            with dst_filesystem.open(save_path, 'ab+') as h5_file_path:
                file_path, patch_time_elapsed = patching(WSI_object, h5_file_path, **current_patch_params,)
        
        stitch_time_elapsed = -1
        if stitch:
            file_path = str(Path(dst_path + patch_save_dir) / f'{slide_id}.h5')
            if dst_filesystem.isfile(file_path):
                with dst_filesystem.open(file_path, 'rb') as hdf5_file_path:
                    heatmap, stitch_time_elapsed = stitching(hdf5_file_path, WSI_object, downscale=64)
                stitch_path = str(Path(dst_path + stitch_save_dir) / f'{slide_id}.jpg')
                with dst_filesystem.open(stitch_path, 'wb') as f:
                    heatmap.save(f, "JPEG")

        print("segmentation took {} seconds".format(seg_time_elapsed))
        print("patching took {} seconds".format(patch_time_elapsed))
        print("stitching took {} seconds".format(stitch_time_elapsed))
        df.loc[idx, 'status'] = 'processed'

        seg_times += seg_time_elapsed
        patch_times += patch_time_elapsed
        stitch_times += stitch_time_elapsed

    seg_times /= total
    patch_times /= total
    stitch_times /= total

    with dst_filesystem.open(csv_path, 'w') as f:
        df.to_csv(f, index=False)
    print("average segmentation time in s per slide: {}".format(seg_times))
    print("average patching time in s per slide: {}".format(patch_times))
    print("average stiching time in s per slide: {}".format(stitch_times))
        
    return seg_times, patch_times

parser = argparse.ArgumentParser(description='seg and patch')
parser.add_argument('--source', type = str,
                    help='path to folder containing raw wsi image files')
parser.add_argument('--step_size', type = int, default=256,
                    help='step_size')
parser.add_argument('--patch_size', type = int, default=256,
                    help='patch_size')
parser.add_argument('--patch', default=False, action='store_true')
parser.add_argument('--seg', default=False, action='store_true')
parser.add_argument('--stitch', default=False, action='store_true')
parser.add_argument('--no_auto_skip', default=True, action='store_false')
parser.add_argument('--save_dir', type = str,
                    help='directory to save processed data')
parser.add_argument('--preset', default=None, type=str,
                    help='predefined profile of default segmentation and filter parameters (.csv)')
parser.add_argument('--patch_level', type=int, default=0, 
                    help='downsample level at which to patch')
parser.add_argument('--process_list',  type = str, default=None,
                    help='name of list of images to process with parameters (.csv)')
parser.add_argument('--s3_storage_key', type = str, default=None)
parser.add_argument('--s3_storage_secret', type = str, default=None)
parser.add_argument('--s3_endpoint_url', type = str, default=None)

if __name__ == '__main__':
    args = parser.parse_args()

    patch_save_dir = '/patches'
    mask_save_dir = '/masks'
    stitch_save_dir = '/stitches'

    if args.process_list:
        process_list = os.path.join(args.save_dir, args.process_list)

    else:
        process_list = None

    print('source: ', args.source)
    print('patch_save_dir: ', patch_save_dir)
    print('mask_save_dir: ', mask_save_dir)
    print('stitch_save_dir: ', stitch_save_dir)
    
    directories = {'source': args.source, 
                   'save_dir': args.save_dir,
                   'patch_save_dir': patch_save_dir, 
                   'mask_save_dir' : mask_save_dir, 
                   'stitch_save_dir': stitch_save_dir} 

    seg_params = {'seg_level': -1, 'sthresh': 8, 'mthresh': 7, 'close': 4, 'use_otsu': False,
                  'keep_ids': 'none', 'exclude_ids': 'none'}
    filter_params = {'a_t':100, 'a_h': 16, 'max_n_holes':8}
    vis_params = {'vis_level': -1, 'line_thickness': 250}
    patch_params = {'use_padding': True, 'contour_fn': 'four_pt'}

    if args.preset:
        preset_df = pd.read_csv(os.path.join('presets', args.preset))
        for key in seg_params.keys():
            seg_params[key] = preset_df.loc[0, key]

        for key in filter_params.keys():
            filter_params[key] = preset_df.loc[0, key]

        for key in vis_params.keys():
            vis_params[key] = preset_df.loc[0, key]

        for key in patch_params.keys():
            patch_params[key] = preset_df.loc[0, key]
    
    parameters = {'seg_params': seg_params,
                  'filter_params': filter_params,
                  'patch_params': patch_params,
                  'vis_params': vis_params}

    print(parameters)

    seg_times, patch_times = seg_and_patch(**directories, **parameters,
                                            patch_size = args.patch_size, step_size=args.step_size, 
                                            seg = args.seg,  use_default_params=False, save_mask = True, 
                                            stitch= args.stitch,
                                            patch_level=args.patch_level, patch = args.patch,
                                            process_list = process_list, auto_skip=args.no_auto_skip,
                                            s3_storage_key=args.s3_storage_key, 
                                            s3_storage_secret=args.s3_storage_secret, 
                                            s3_endpoint_url=args.s3_endpoint_url)