This repo contains three ONNX variants exported from the official PyTorch checkpoints, covering both low-light enhancement and exposure correction tasks.
1import numpy as np
2from PIL import Image
3
4img = Image.open("dark_photo.jpg").convert("RGB")
5img_np = np.array(img).astype(np.float32) / 255.0 # [0, 1]
6# Transpose to CHW and add batch dim
7input_tensor = img_np.transpose(2, 0, 1)[np.newaxis, ...] # (1, 3, H, W)
1import numpy as np
2import onnxruntime as ort
3from PIL import Image
4
5# Load model
6session = ort.InferenceSession("onnx/iat_lol_v2.onnx", providers=["CPUExecutionProvider"])
7
8# Preprocess
9img = Image.open("dark_photo.jpg").convert("RGB")
10img_np = np.array(img).astype(np.float32) / 255.0
11input_tensor = img_np.transpose(2, 0, 1)[np.newaxis, ...] # (1, 3, H, W)
12
13# Run inference — use "enhanced" (index 2)
14mul, add, enhanced = session.run(None, {"input": input_tensor})
15
16# Post-process
17enhanced = np.clip(enhanced[0], 0, 1) # (3, H, W)
18enhanced = (enhanced.transpose(1, 2, 0) * 255).astype(np.uint8) # (H, W, 3)
19result = Image.fromarray(enhanced)
20result.save("enhanced.jpg")
-
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.
-
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.
-
Aff_channel.forward: Same tensordot to matmul fix as patch 1, applied to the channel affinity block in the local branch.
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).
1@InProceedings{Cui_2022_BMVC,
2 title = {Illumination Adaptive Transformer},
3 author = {Cui, Ziteng and Li, Kunchang and Gu, Lin and Su, Shenghan and Gao, Peng and Jiang, Zhengkai and Qiao, Yu and Harada, Tatsuya},
4 booktitle = {British Machine Vision Conference (BMVC)},
5 year = {2022}
6}
Apache-2.0 — same as the original IAT repository.
