COSIE.image_preprocessing.image_feature_extraction

image_feature_extraction(he_image, uni_local_dir, cell_location, device='cuda:0', batch_size=128, num_workers=4)

Extract image features at cell locations using a pretrained Vision Transformer (ViT) model and save the output to disk as a pickle file.

The function loads a model via create_model, initializes a PatchDataset using the input image and cell pixel locations, and for each patch extracts both global (224×224) and local (16×16) features. These features are concatenated into a single representation per cell.

Parameters

he_imagenp.ndarray

RGB image of shape (H, W, 3).

uni_local_dirstr

Path to the folder containing the pretrained model weights (e.g., pytorch_model.bin).

cell_locationnp.ndarray

An array of shape (N, 2) containing pixel coordinates for N cells.

devicestr, optional

Torch device to use for inference, e.g., “cuda:0”. Default is ‘cuda:0’.

batch_sizeint, optional

Batch size for feature extraction. Default is 128.

num_workersint, optional

Number of worker threads for the DataLoader. Default is 4.

Returns

None

The extracted features are saved to disk as a pickle file.