litert-community/DehazeFormer-MCT-LiteRT

DehazeFormer-MCT — Image dehazing (LiteRT GPU)

On-device image dehazing with the network fully on the LiteRT CompiledModel GPU delegate (no CPU fallback). DehazeFormer (TIP 2023, MCT curve-mapping variant, trained by the author on a mixed dataset for real-world haze) removes fog / haze / smoke and restores contrast and color.

• Architecture: DehazeFormer basenet (Swin-style windowed attention, 1.2M params) → 72 per-pixel curve parameters.
• Weights: author-hosted IDKiro/DehazeFormer_Demo · MIT.
• Size: 17 MB.

Hazy input (left) → dehazed (right). Photo: Pexels (free license).

The MCT design is mobile-ideal: the network always runs at 256×256; the predicted per-pixel curves are applied to the full-resolution image host-side (a cheap trilinear lookup — the official grid_sample mapping), so output resolution is independent of the network.

I/O

• Input: [1, 3, 256, 256] NCHW, RGB in [-1, 1] (x/255*2-1).
• Output: [1, 72, 256, 256] curve parameters — layout [3 out-channels × 3 in-channels × 8 levels].
• Host mapping (per full-res pixel): out[c] = Σᵢ trilinear(curve[c][i], depth = xᵢ, y, x) with align_corners=true and border clamping, then clamp(-1,1)*0.5+0.5.

GPU conversion

Fully GPU-resident on a Pixel 8a (2042/2042 nodes, 1 partition; device corr 0.999998, end-to-end vs the official pipeline corr 0.999997, ~255 ms/frame) via exact re-authors: reflect pads → slice+concat (litert-torch lowers reflection_pad2d to GATHER_ND, rejected by the delegate), Swin window partition/reverse in ≤4D + baked relative-position bias, SKFusion 5D→4D pairwise softmax, Conv+PixelShuffle → zero-stuff ConvTranspose, and — the new finding — hierarchical means for the RLN global norm (a single MEAN over 1.5M elements overflows the Mali fp16 accumulator → NaN; equal-window avg_pool stages are mathematically identical and fp16-safe). Desktop corr vs PyTorch is 1.0000000.

Minimal usage

Kotlin (Android, LiteRT CompiledModel GPU)

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

inBufs[0].writeFloat(inputNCHW)          // [1,3,256,256] RGB in [-1,1]
model.run(inBufs, outBufs)
val curves = outBufs[0].readFloat()      // [72*256*256] curve params
// apply curves to the full-res frame host-side (see the sample's Dehazer.applyCurves)

Python (LiteRT CompiledModel API)

import numpy as np
from ai_edge_litert.compiled_model import CompiledModel

model = CompiledModel.from_file("dehazeformer_base.tflite")
inputs = model.create_input_buffers(0)
outputs = model.create_output_buffers(0)
inputs[0].write(np.ascontiguousarray(x, np.float32))  # [1,3,256,256] RGB in [-1,1]
model.run_by_index(0, inputs, outputs)
n = model.get_output_buffer_requirements(0, 0)["buffer_size"] // 4
curves = outputs[0].read(n, np.float32).reshape(72, 256, 256)

Conversion

Converted with litert-torch (build_dehaze.py): fetches the author's model code and MIT checkpoint from the Hugging Face Space and exports the curve-parameter basenet.

License

MIT (DehazeFormer / IDKiro). Mixed-dataset checkpoint from the author's demo Space.