CanViT COCO detector (JAX / Flax NNX)

CanViT-B with a CenterNet head on its canvas, trained end to end for COCO 2017 object detection; it predicts boxes after every glimpse.

CanViT, the Canvas Vision Transformer, is an active-vision foundation model: it sees a scene through a sequence of glimpses and remembers it on a scene-wide canvas.

Paper (NeurIPS 2026) · Code · Project page · All checkpoints

Install

uv add canvit-nnx

Usage

import flax.nnx as nnx
import jax
import jax.numpy as jnp
import numpy as np
from PIL import Image

from canvit_nnx import sample_at_viewpoint
from canvit_nnx.detect.data import normalize
from canvit_nnx.detect.decode import decode
from canvit_nnx.detect.model import CanViTForObjectDetection
from canvit_nnx.policies import coarse_to_fine

image = Image.open("scene.jpg").convert("RGB").resize((512, 512), Image.Resampling.BILINEAR)
scene = normalize(jnp.asarray(np.asarray(image))[None])
model = nnx.view(CanViTForObjectDetection.from_pretrained("canvit/canvitb16-add-vpe-detect-g128px-s512px-coco-from-archctrl-2026-07-27-nnx"), use_running_average=True)
state = model.init_state(batch_size=1, canvas_grid_size=32)
for viewpoint in coarse_to_fine(jax.random.key(0), batch_size=1, num_glimpses=5):
    glimpse = sample_at_viewpoint(spatial=scene, viewpoint=viewpoint, glimpse_size_px=128)
    maps, state = model(glimpse=glimpse, state=state, viewpoint=viewpoint)
    detections = decode(maps)  # normalized (ymin, xmin, ymax, xmax) boxes, scores and COCO class indices

Results

COCO val2017 (5000 images), C2F viewpoints: means over 5 independent viewpoint draws, with 95% bootstrap intervals over the draws where they vary.

The paper prints this AP, in percent, at glimpses 0, 1, 2, 3 and 4, in the row "CanViT-B (architecture control)" of its table "COCO val2017 AP by timestep, C2F policy" (appendix "COCO object detection"). python -m canvit_nnx.detect.evaluate (canvit_nnx.detect) scores a detector on val2017; it needs canvit-nnx[train].

Details

License

MIT. CanViT is pretrained by distillation from DINOv3, which is released under the DINOv3 License.

Citation

@inproceedings{berreby2026canvit,
  title={CanViT: Toward Active-Vision Foundation Models},
  author={Berreby, Yoha{\"i}-Eliel and Du, Sabrina and Durand, Audrey and Krishna, B. Suresh},
  booktitle={Advances in Neural Information Processing Systems},
  year={2026},
  eprint={2603.22570},
  archivePrefix={arXiv},
  primaryClass={cs.CV},
  url={https://arxiv.org/abs/2603.22570}
}