tensorkelechi/DReX-text-tuned

DReX, tuned for text

The last layers of DReX, retrained so that pages of text no longer score as simple images.

Original model: code and DReX.pth weights, paper. You need the original checkpoint, because this repo holds only the replaced layers.

What DReX does

DReX predicts how visually complex people find an image, as a single score. It combines DINOv3 and ResNet-50 features and was trained on human ratings from the IC9600 dataset.

What this version was tuned for

People rate a page of text as fairly plain, so DReX gives it a low score. We use the score to decide how many tokens an image gets in a variable-length image tokenizer, and there a low score makes text unreadable. This version scores text by the tokens it needs to stay legible.

Only the fusion and head layers changed, and the backbones are frozen. The training labels came from OCR on document images from DocLayNet, FUNSD and SynthDoG, and a penalty held non-text images to their original score.

This is no longer a pure predictor of human complexity ratings. On text it predicts what a tokenizer needs. If you want the human rating, use the original.

Results

In our tokenizer on 63 held-out pages, counting a page as readable when IndraOCR recovers at least 90% of the words it reads on the original:

Budget from	Pages readable	Median tokens
Original DReX	8%	784
This version	60%	2,304
Full grid, no adaptation	65%	2,304

On 52 test images, nearly all non-document, the median score moved by .015, with a Spearman correlation of .995 against the original. Median Mondrian budgets rose from 60 to 75 tokens.

Limits: tested on printed documents only, not scene text or receipts. It missed the 80% readable bar we set beforehand, a bar the full grid itself does not reach.

Use

import torch
from huggingface_hub import hf_hub_download

w = torch.load(hf_hub_download("tensorkelechi/DReX-text-tuned", "drext.pt"), weights_only=True)
model.fusion.load_state_dict(w["fusion"])
model.head.load_state_dict(w["head"])

model is the original DReX, already loaded from DReX.pth.

License

MIT, as the original DReX. Copyright (c) 2025 Jonathan Skaza; see LICENSE.