LESSViT-Base / load_encoder.py

Use public model name LESSViT-Base in documentation and examples

bbb9d1f verified 1 day ago

1.8 kB

"""Load the released encoder and extract features from a synthetic HSI batch."""
import argparse
import torch
from transformers import AutoModel


def load_encoder(model_id="GFM-Bench/LESSViT-Base", device="cpu", revision=None):
    model = AutoModel.from_pretrained(
        model_id, trust_remote_code=True, revision=revision,
    )
    return model.to(device).eval()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model-id', default='GFM-Bench/LESSViT-Base',
                        help='Hugging Face model ID or local snapshot directory')
    parser.add_argument('--revision', default=None, help='Optional commit SHA')
    parser.add_argument('--device', default='cpu')
    args = parser.parse_args()
    torch.set_num_threads(4)
    encoder = load_encoder(args.model_id, args.device, args.revision)
    # Synthetic values demonstrate tensor shapes, not preprocessing for real imagery.
    image = torch.randn(1, 10, 32, 32, device=args.device)
    wavelengths_nm = torch.linspace(450, 2300, 10, device=args.device)[None]
    # Example: fixed full-EnMAP wavelength bounds. Use the same reference bounds
    # for train/eval in your downstream pipeline; do not renormalize each subset.
    wavelengths = (wavelengths_nm - 418.24) / (2445.53 - 418.24)
    with torch.no_grad():
        output = encoder(optical=image, optical_channel_wv=wavelengths,
                         spatial_resolution=30)
    grid = output.last_hidden_state
    assert torch.isfinite(grid).all()
    print('Token grid:', tuple(grid.shape))
    print('Global features:', tuple(grid[:, 0, 0].shape))
    print('Spatial features:', tuple(grid[:, 0, 1:].shape))
    print('Per-band patch features:', tuple(grid[:, 1:, 1:].shape))


if __name__ == '__main__':
    main()

	"""Load the released encoder and extract features from a synthetic HSI batch."""
	import argparse
	import torch
	from transformers import AutoModel


	def load_encoder(model_id="GFM-Bench/LESSViT-Base", device="cpu", revision=None):
	model = AutoModel.from_pretrained(
	model_id, trust_remote_code=True, revision=revision,
	)
	return model.to(device).eval()


	def main():
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument('--model-id', default='GFM-Bench/LESSViT-Base',
	help='Hugging Face model ID or local snapshot directory')
	parser.add_argument('--revision', default=None, help='Optional commit SHA')
	parser.add_argument('--device', default='cpu')
	args = parser.parse_args()
	torch.set_num_threads(4)
	encoder = load_encoder(args.model_id, args.device, args.revision)
	# Synthetic values demonstrate tensor shapes, not preprocessing for real imagery.
	image = torch.randn(1, 10, 32, 32, device=args.device)
	wavelengths_nm = torch.linspace(450, 2300, 10, device=args.device)[None]
	# Example: fixed full-EnMAP wavelength bounds. Use the same reference bounds
	# for train/eval in your downstream pipeline; do not renormalize each subset.
	wavelengths = (wavelengths_nm - 418.24) / (2445.53 - 418.24)
	with torch.no_grad():
	output = encoder(optical=image, optical_channel_wv=wavelengths,
	spatial_resolution=30)
	grid = output.last_hidden_state
	assert torch.isfinite(grid).all()
	print('Token grid:', tuple(grid.shape))
	print('Global features:', tuple(grid[:, 0, 0].shape))
	print('Spatial features:', tuple(grid[:, 0, 1:].shape))
	print('Per-band patch features:', tuple(grid[:, 1:, 1:].shape))


	if __name__ == '__main__':
	main()