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.

License

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