Function extract_features

extract_features.py:555–574

(model, device, wsi, filtered_tiles, workers, out_size, batch_size, n_last_blocks, avgpool_patchtokens, depths)

@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.