scitex_ml.feature_extraction.vit

Module Attributes

THIS_FILE

Functionality:

Classes

VitFeatureExtractor([model_name, ...])

scitex_ml.feature_extraction.vit.THIS_FILE = '/home/ywatanabe/proj/scitex_repo/src/scitex/ai/feature_extraction/vit.py'
Functionality:

Extracts features from images using Vision Transformer (ViT) models

Input:

Image arrays of arbitrary dimensions

Output:

Feature vectors (1000-dimensional embeddings)

Prerequisites:

torch, PIL, torchvision

class scitex_ml.feature_extraction.vit.VitFeatureExtractor(model_name='B_16', torch_home='./models', device=None)[source]
__init__(model_name='B_16', torch_home='./models', device=None)[source]
extract_features(arr, axis=(-2, -1), dim=None, channel_dim=None, batch_size=None, device='cuda')[source]