--- a
+++ b/extract_features.py
@@ -0,0 +1,652 @@
+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)
\ No newline at end of file