UHR-BAT is a budget-aware vision-language framework for ultra-high-resolution remote sensing imagery. It targets the setting where kilometer-scale scenes contain query-critical evidence that may occupy only a few pixels. Instead of relying on direct downsampling, dense tiling, or generic global pruning, UHR-BAT uses query-guided multi-scale token selection and region-faithful compression to preserve small decisive evidence under a strict context budget.
The project page reports strong ultra-high-resolution remote-sensing results under strict token budgets:
1import importlib
2import torch
3from PIL import Image
4from transformers import AutoImageProcessor
5from transformers import AutoModelForCausalLM, AutoTokenizer
6
7model_id = "FelixKAI/UHR-BAT"
8image_path = "your_remote_sensing_image.jpg"
9question = "Describe this remote-sensing image."
10
11tokenizer = AutoTokenizer.from_pretrained(model_id)
12if tokenizer.pad_token_id is None:
13 tokenizer.pad_token = tokenizer.eos_token
14image_processor = AutoImageProcessor.from_pretrained(model_id)
15
16model = AutoModelForCausalLM.from_pretrained(
17 model_id,
18 trust_remote_code=True,
19 torch_dtype="auto",
20 device_map="auto",
21).eval()
22
23# Reuse the preprocessing helpers shipped with the model's remote code.
24uhrbat = importlib.import_module(model.__class__.__module__)
25image = Image.open(image_path).convert("RGB")
26
27prompt = (
28 "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
29 f"<|im_start|>user\n<image>\n{question}<|im_end|>\n"
30 "<|im_start|>assistant\n"
31)
32image_token_id = getattr(model.config, "image_token_index", -200)
33input_ids = uhrbat.tokenizer_image_token(
34 prompt,
35 tokenizer,
36 image_token_id,
37 return_tensors="pt",
38).unsqueeze(0).to(model.device)
39attention_mask = torch.ones_like(input_ids)
40
41target_sizes = [672, 1344, 2688, 4032]
42multiscale_pixels = [
43 uhrbat.split_image_to_multiscale_tiles(
44 image,
45 image_processor,
46 target_sizes=target_sizes,
47 tile_size=336,
48 )
49]
50
51with torch.inference_mode():
52 output = model.generate(
53 inputs=input_ids,
54 attention_mask=attention_mask,
55 image_sizes=[image.size],
56 modalities=["image"],
57 multiscale_pixels=multiscale_pixels,
58 multiscale_masks=[{}],
59 multiscale_topk=[80, 320, 600, 2000],
60 multiscale_target_sizes=target_sizes,
61 do_sample=False,
62 max_new_tokens=256,
63 return_dict_in_generate=True,
64 output_scores=True,
65 pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
66 eos_token_id=tokenizer.eos_token_id,
67 )
68
69prompt_len = output.sequences.shape[1] - len(output.scores)
70answer_ids = output.sequences[:, prompt_len:].clone()
71answer_ids[answer_ids < 0] = tokenizer.pad_token_id or tokenizer.eos_token_id
72answer = tokenizer.decode(answer_ids[0], skip_special_tokens=True).strip()
73print(answer)
1@inproceedings{dang2026uhrbat,
2 title={UHR-BAT: Budget-Aware Token Compression Vision-Language model for Ultra-High-Resolution Remote Sensing},
3 author={Dang, Yunkai and Dai, Minxin and Yang, Yuekun and Li, Zhangnan and Li, Wenbin and Miao, Feng and Gao, Yang},
4 booktitle={International Conference on Machine Learning (ICML)},
5 year={2026}
6}
