IAT-ONNX

IAT — Illumination Adaptive Transformer (ONNX)

First public ONNX export of the Illumination Adaptive Transformer (IAT) by Cui et al.

This repo contains three ONNX variants exported from the official PyTorch checkpoints, covering both low-light enhancement and exposure correction tasks.

Variants

Model Specs

Preprocessing

import numpy as np
from PIL import Image

img = Image.open("dark_photo.jpg").convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0   # [0, 1]
# Transpose to CHW and add batch dim
input_tensor = img_np.transpose(2, 0, 1)[np.newaxis, ...]  # (1, 3, H, W)

Important: Do NOT apply ImageNet normalization. The model expects raw [0, 1] pixel values.

Usage with ONNX Runtime

import numpy as np
import onnxruntime as ort
from PIL import Image

# Load model
session = ort.InferenceSession("onnx/iat_lol_v2.onnx", providers=["CPUExecutionProvider"])

# Preprocess
img = Image.open("dark_photo.jpg").convert("RGB")
img_np = np.array(img).astype(np.float32) / 255.0
input_tensor = img_np.transpose(2, 0, 1)[np.newaxis, ...]  # (1, 3, H, W)

# Run inference — use "enhanced" (index 2)
mul, add, enhanced = session.run(None, {"input": input_tensor})

# Post-process
enhanced = np.clip(enhanced[0], 0, 1)               # (3, H, W)
enhanced = (enhanced.transpose(1, 2, 0) * 255).astype(np.uint8)  # (H, W, 3)
result = Image.fromarray(enhanced)
result.save("enhanced.jpg")

ONNX Export Fixes

The original PyTorch code required three monkey-patches for clean ONNX tracing:

1. IAT.apply_color: Replaced torch.tensordot(image, ccm, dims=[[-1], [-1]]) with torch.matmul(image, ccm.T) — tensordot with negative dimension indices is not supported by the ONNX exporter.

2. IAT.forward: Replaced Python for-loop over the batch dimension (for i in range(b)) with vectorized torch.bmm for the color matrix multiply and broadcast ** for gamma correction. Python loops produce unrollable static graphs that break with dynamic batch sizes.

3. Aff_channel.forward: Same tensordot to matmul fix as patch 1, applied to the channel affinity block in the local branch.

See export_iat_onnx.py in this repo for the full export script with patches.

Architecture

IAT is a lightweight image enhancement model with two branches:

• Local branch: Learns per-pixel multiplicative (mul) and additive (add) adjustment maps via a shallow transformer. enhanced_local = input * mul + add
• Global branch: Learns a 3x3 color correction matrix (CCM) and a scalar gamma value. Applied after local enhancement: enhanced = (enhanced_local @ CCM^T) ^ gamma

The combination of local pixel-wise adjustments and global color/tone correction makes it effective for both low-light enhancement and exposure correction, while keeping the model extremely small (~90K parameters).

Benchmark Results

Results from the original paper:

Citation

@InProceedings{Cui_2022_BMVC,
    title     = {Illumination Adaptive Transformer},
    author    = {Cui, Ziteng and Li, Kunchang and Gu, Lin and Su, Shenghan and Gao, Peng and Jiang, Zhengkai and Qiao, Yu and Harada, Tatsuya},
    booktitle = {British Machine Vision Conference (BMVC)},
    year      = {2022}
}

License

Apache-2.0 — same as the original IAT repository.

Acknowledgments

• Original model and research by Cui et al.
• ONNX export by ListingGems — an offline product photo editor for marketplace sellers