litert-community/DewarpNet-LiteRT

DewarpNet — Document unwarping (LiteRT GPU)

On-device document dewarping / rectification running fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). DewarpNet (ICCV 2019) flattens a photographed, curved/folded document — the core of a document scanner. Two CNNs predict a backward-mapping grid; the network runs on the GPU and the grid_sample unwarp is a tiny host-side step. ~24 ms/frame on a Pixel 8a.

• Architecture: WCNet (UNet, world-coords) → BMNet (DenseNet, backward map) — pure CNN.
• Weights: cvlab-stonybrook/DewarpNet (doc3d) · MIT.
• Size: 189 MB.

Left: photographed curved page. Right: dewarped/rectified. Input photo: Unsplash (free license).

I/O

• Input: [1, 3, 256, 256] NCHW, BGR, x/255.
• Output: [1, 2, 128, 128] backward-mapping grid (values ~`[-1,1]`).
• Host-side unwarp: blur the map (3×3), resize to the original image size, then grid_sample(original_image, map) → the flattened document.

GPU conversion

DewarpNet is a pure CNN. It converts fully GPU-compatible (371/371 nodes on the delegate, 1 partition; device corr 0.999866, ~24 ms) with two patches: (1) the UNet/DenseNet ConvTranspose2d upsamplers → ZeroStuffConvT2d (nearest-upsample + stride zero-stuff mask + flipped conv; the Mali delegate rejects TRANSPOSE_CONV); and (2) Hardtanh(0,1) → relu(x) - relu(x-1) (the delegate rejects RELU_0_TO_1). Both are exact. CPU-exact vs PyTorch (corr 0.9999999999).

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

val options = CompiledModel.Options(Accelerator.GPU)
val model = CompiledModel.create(context.assets, "dewarp.tflite", options, null)
val inBufs = model.createInputBuffers()
val outBufs = model.createOutputBuffers()

inBufs[0].writeFloat(inputNCHW)          // [1,3,256,256] BGR, x/255
model.run(inBufs, outBufs)
val bm = outBufs[0].readFloat()          // [2*128*128] backward map (grid, ~[-1,1])
// host: blur 3x3, resize to image size, then bilinear grid_sample(image, bm) -> flattened doc

Python (LiteRT / ai-edge-litert)

import numpy as np, cv2, torch, torch.nn.functional as F
from ai_edge_litert.interpreter import Interpreter

it = Interpreter(model_path="dewarp.tflite"); it.allocate_tensors()
inp, out = it.get_input_details(), it.get_output_details()
it.set_tensor(inp[0]["index"], x)        # [1,3,256,256] float32, BGR, x/255
it.invoke()
bm = it.get_tensor(out[0]["index"])      # [1,2,128,128]
bm = np.stack([cv2.resize(cv2.blur(bm[0,0],(3,3)), (W,H)),
               cv2.resize(cv2.blur(bm[0,1],(3,3)), (W,H))], -1)[None]
flat = F.grid_sample(torch.tensor(imgorg/255.).permute(2,0,1)[None].float(),
                     torch.tensor(bm).float(), align_corners=True)   # unwarped

Conversion

Converted with litert-torch (build_dewarp.py): loads the two CNNs, applies the ZeroStuffConvT2d + clamp patches, and exports the image→backward-map graph.

Performance

Measured on a Pixel 8a (Tensor G3, Android 16) with the standard TFLite benchmark_model tool — 10 warm-up runs then 50 timed runs, reported as the tool's mean.

Runtime	Backend	Graph on GPU	Latency
LiteRT CompiledModel (LITERT_CL)	GPU	371 / 371	~24 ms
TFLite benchmark_model (TfLiteGpuDelegateV2)	GPU (OpenCL)	371 / 371	88.3 ms
TFLite benchmark_model	CPU (XNNPACK, 4 threads)	—	875.2 ms

The two GPU rows are different runtimes, not a contradiction. The LITERT_CL figure is the one recorded when this model shipped, taken through LiteRT's own CompiledModel accelerator — the path the Kotlin sample app and the LiteRT API use. The TfLiteGpuDelegateV2 figure is the classic TFLite OpenCL delegate, measured with a tool anyone can download and re-run. They agree on how much of the graph the GPU takes; they disagree on speed, and the classic delegate is the slower of the two here. Read the TfLiteGpuDelegateV2 row as a reproducible floor, not as this model's speed on LiteRT.

Snapdragon NPU (Hexagon)

The NPU is 3.53x faster than the GPU (4.99 ms against 17.64 ms) and loads 6.25x faster (199 ms against 1243 ms).

backend	inference (median / min)	load
NPU (Hexagon v81)	4.99 ms / 4.95 ms	199 ms
GPU (Adreno)	17.64 ms / 17.38 ms	1243 ms

Measured on a Samsung Galaxy S26 (Snapdragon 8 Elite Gen 5 / SM8850, Hexagon v81, Android 16), LiteRT CompiledModel 2.2.0, one accelerator per process, 5 warm-up runs then N=50 timed runs, median reported. Every run held thermal status NONE throughout. Headroom 0.71-0.74, where 1.0 is the throttling threshold.

The NPU rows here ran artifacts compiled ahead of time for SM8850 with QAIRT 2.47.0; the GPU rows ran the published files as they are. LiteRT can also compile for the NPU on the device at first load, which is what lets you ship the published file unchanged — that path and the ten runtime libraries it needs are in the NPU recipe, and we did not measure it here. GPU wiring is in the GPU recipe.

License

MIT (DewarpNet / cvlab-stonybrook). Trained on the Doc3D dataset.