Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+onnx/iat_exposure.onnx.data filter=lfs diff=lfs merge=lfs -text
+onnx/iat_lol_v1.onnx.data filter=lfs diff=lfs merge=lfs -text
+onnx/iat_lol_v2.onnx.data filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,126 @@
+---
+license: apache-2.0
+base_model: cuiziteng/Illumination-Adaptive-Transformer
+pipeline_tag: image-to-image
+tags:
+  - low-light-enhancement
+  - exposure-correction
+  - image-enhancement
+  - onnx
+---
+
+# IAT — Illumination Adaptive Transformer (ONNX)
+
+**First public ONNX export** of the [Illumination Adaptive Transformer (IAT)](https://github.com/cuiziteng/Illumination-Adaptive-Transformer) 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
+
+| File | Checkpoint | Training Data | Use Case |
+|------|-----------|---------------|----------|
+| `onnx/iat_exposure.onnx` | `best_Epoch_exposure.pth` | [Exposure Errors](https://github.com/mahmoudnafifi/Exposure_Correction) | Over/under-exposure correction |
+| `onnx/iat_lol_v1.onnx` | `best_Epoch_lol_v1.pth` | [LOL-V1](https://daooshee.github.io/BMVC2018website/) | Low-light enhancement |
+| `onnx/iat_lol_v2.onnx` | `best_Epoch_lol.pth` | [LOL-V2](https://github.com/flyywh/CVPR-2020-Semi-Low-Light) | Low-light enhancement (improved) |
+
+## Model Specs
+
+| Property | Value |
+|----------|-------|
+| Parameters | ~90K |
+| File size | ~0.1 MB per variant |
+| Input shape | `(1, 3, H, W)` float32, values in `[0, 1]` |
+| Normalization | **None** — just rescale to [0,1], no ImageNet mean/std |
+| Output names | `mul`, `add`, `enhanced` |
+| Which output to use | `enhanced` (index 2) |
+| Dynamic axes | batch, height, width |
+| ONNX opset | 17 |
+
+## Preprocessing
+
+```python
+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
+
+```python
+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:
+
+| Dataset | PSNR | SSIM |
+|---------|------|------|
+| LOL-V1 | 23.38 | 0.809 |
+| LOL-V2 | 23.50 | 0.824 |
+
+## Citation
+
+```bibtex
+@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.](https://github.com/cuiziteng/Illumination-Adaptive-Transformer)
+- ONNX export and this model card by [ListingLens](https://github.com/Pezhgorski/listinglens)

config.json

@@ -0,0 +1,3 @@
+{
+  "model_type": "iat"
+}

export_iat_onnx.py

@@ -0,0 +1,372 @@
+#!/usr/bin/env python3
+"""
+IAT (Illumination Adaptive Transformer) → ONNX Export Script
+
+Monkey-patches ONNX-incompatible patterns in IAT source, exports all 3
+checkpoints (exposure, lol_v1, lol_v2), and verifies each numerically
+at multiple resolutions.
+
+Patches applied:
+  1. IAT.apply_color: tensordot → matmul (ONNX-friendly)
+  2. IAT.forward: Python for-loop over batch → vectorized bmm + broadcast pow
+  3. Aff_channel.forward: tensordot → matmul (fallback — needed for tracing)
+"""
+
+import argparse
+import sys
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+# ---------------------------------------------------------------------------
+# Add IAT source to path, fix Python 3.12+ compatibility
+# ---------------------------------------------------------------------------
+IAT_ROOT = Path(__file__).parent / "iat" / "IAT_enhance"
+sys.path.insert(0, str(IAT_ROOT))
+
+# IAT's global_net.py has `import imp` which was removed in Python 3.12.
+# It's unused, so we provide a dummy module before importing IAT.
+import importlib
+if not importlib.util.find_spec("imp"):
+    import types
+    sys.modules["imp"] = types.ModuleType("imp")
+
+from model.IAT_main import IAT        # noqa: E402
+from model.blocks import Aff_channel   # noqa: E402
+
+
+# ===========================================================================
+# Monkey-patches
+# ===========================================================================
+
+def _patched_apply_color(self, image, ccm):
+    """Replace tensordot with matmul for ONNX compatibility.
+
+    Original: torch.tensordot(image, ccm, dims=[[-1], [-1]])
+      which computes image @ ccm.T  (contract last dim of both)
+    Replacement: torch.matmul(image, ccm.T)
+    """
+    shape = image.shape
+    image = image.view(-1, 3)
+    image = torch.matmul(image, ccm.permute(1, 0))
+    image = image.view(shape)
+    return torch.clamp(image, 1e-8, 1.0)
+
+
+def _patched_forward(self, img_low):
+    """Vectorized forward — no Python for-loop over batch dimension.
+
+    Original:
+        img_high = torch.stack([self.apply_color(img_high[i,:,:,:], color[i,:,:])
+                                **gamma[i,:] for i in range(b)], dim=0)
+
+    Replacement:
+        1. bmm for batched color matrix multiply
+        2. broadcast pow for gamma
+    """
+    mul, add = self.local_net(img_low)
+    img_high = (img_low.mul(mul)).add(add)
+
+    if not self.with_global:
+        return mul, add, img_high
+
+    gamma, color = self.global_net(img_low)
+    # img_high: (B, C, H, W) → (B, H, W, C) → (B, H*W, C)
+    b, c, h, w = img_high.shape
+    img_high = img_high.permute(0, 2, 3, 1).reshape(b, h * w, c)
+
+    # Batched color matrix: (B, H*W, 3) @ (B, 3, 3) → (B, H*W, 3)
+    # color is (B, 3, 3), we need img @ color^T for each batch element
+    color_t = color.permute(0, 2, 1)  # (B, 3, 3)
+    img_high = torch.bmm(img_high, color_t)
+    img_high = torch.clamp(img_high, 1e-8, 1.0)
+
+    # Reshape back to (B, H, W, C) for broadcast pow
+    img_high = img_high.view(b, h, w, c)
+
+    # gamma is (B, 1) — reshape to (B, 1, 1, 1) for broadcast
+    gamma_broadcast = gamma.unsqueeze(-1).unsqueeze(-1)  # (B, 1, 1, 1)
+    img_high = img_high ** gamma_broadcast
+
+    # (B, H, W, C) → (B, C, H, W)
+    img_high = img_high.permute(0, 3, 1, 2)
+    return mul, add, img_high
+
+
+def _patched_aff_channel_forward(self, x):
+    """Replace tensordot with matmul in Aff_channel for ONNX compatibility.
+
+    Original: torch.tensordot(x, self.color, dims=[[-1], [-1]])
+    Replacement: torch.matmul(x, self.color.T)
+    """
+    if self.channel_first:
+        x1 = torch.matmul(x, self.color.permute(1, 0))
+        x2 = x1 * self.alpha + self.beta
+    else:
+        x1 = x * self.alpha + self.beta
+        x2 = torch.matmul(x1, self.color.permute(1, 0))
+    return x2
+
+
+# ===========================================================================
+# Fallback patches (not needed for current export, documented for reference)
+# ===========================================================================
+
+# --- Fallback: query_Attention expand ---
+# If export fails on expand in global attention (global_net.py):
+#
+# def _patched_query_attention_forward(self, x):
+#     B, N, C = x.shape
+#     # Original: self.q.expand(B, -1, -1) -- can fail with dynamic batch
+#     # Fix: use repeat which traces cleanly
+#     q = self.q.repeat(B, 1, 1).view(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+#     ... rest of forward unchanged ...
+#
+# from model.global_net import query_Attention
+# query_Attention.forward = _patched_query_attention_forward
+
+# --- Fallback: gamma power operator ---
+# If ** operator traces incorrectly for broadcast shapes:
+#
+# Replace in _patched_forward:
+#   img_high = torch.pow(torch.clamp(img_high, 1e-8), gamma)
+# With:
+#   img_high = torch.exp(torch.log(torch.clamp(img_high, 1e-8)) * gamma)
+
+
+# ===========================================================================
+# Checkpoint configurations
+# ===========================================================================
+
+CHECKPOINTS = {
+    "exposure": {
+        "path": IAT_ROOT / "best_Epoch_exposure.pth",
+        "model_kwargs": {"type": "exp"},
+        "description": "Exposure correction",
+    },
+    "lol_v1": {
+        "path": IAT_ROOT / "best_Epoch_lol_v1.pth",
+        "model_kwargs": {"type": "lol"},
+        "description": "LOL v1 low-light enhancement",
+    },
+    "lol_v2": {
+        "path": IAT_ROOT / "best_Epoch_lol.pth",
+        "model_kwargs": {"type": "lol"},
+        "description": "LOL v2 low-light enhancement",
+    },
+}
+
+VERIFICATION_RESOLUTIONS = [
+    (256, 256),
+    (512, 512),
+    (768, 1024),  # H, W — non-square
+]
+
+
+# ===========================================================================
+# Apply patches
+# ===========================================================================
+
+def apply_patches():
+    """Monkey-patch IAT classes at runtime. Does not modify source files."""
+    # Patch 1 & 2: IAT.apply_color and IAT.forward
+    IAT.apply_color = _patched_apply_color
+    IAT.forward = _patched_forward
+
+    # Patch 3 (fallback — needed for Aff_channel tensordot tracing):
+    Aff_channel.forward = _patched_aff_channel_forward
+
+    print("[PATCH] IAT.apply_color: tensordot -> matmul")
+    print("[PATCH] IAT.forward: for-loop -> vectorized bmm + broadcast pow")
+    print("[PATCH] Aff_channel.forward: tensordot -> matmul")
+
+
+# ===========================================================================
+# Export
+# ===========================================================================
+
+def load_model(name: str) -> nn.Module:
+    """Load an IAT model from checkpoint."""
+    cfg = CHECKPOINTS[name]
+    model = IAT(in_dim=3, with_global=True, **cfg["model_kwargs"])
+    state_dict = torch.load(str(cfg["path"]), map_location="cpu", weights_only=True)
+    model.load_state_dict(state_dict)
+    model.train(False)
+    return model
+
+
+def export_onnx(model: nn.Module, output_path: Path, opset: int) -> None:
+    """Export a single IAT model to ONNX."""
+    dummy_input = torch.randn(1, 3, 256, 256)
+
+    torch.onnx.export(
+        model,
+        (dummy_input,),
+        str(output_path),
+        opset_version=opset,
+        input_names=["input"],
+        output_names=["mul", "add", "enhanced"],
+        dynamic_axes={
+            "input": {0: "batch", 2: "height", 3: "width"},
+            "mul": {0: "batch", 2: "height", 3: "width"},
+            "add": {0: "batch", 2: "height", 3: "width"},
+            "enhanced": {0: "batch", 2: "height", 3: "width"},
+        },
+    )
+
+
+def verify_onnx(model: nn.Module, onnx_path: Path) -> bool:
+    """Numerical verification of ONNX vs PyTorch at multiple resolutions."""
+    import onnxruntime as ort
+
+    session = ort.InferenceSession(
+        str(onnx_path),
+        providers=["CPUExecutionProvider"],
+    )
+
+    all_ok = True
+    for h, w in VERIFICATION_RESOLUTIONS:
+        dummy = torch.randn(1, 3, h, w)
+
+        # PyTorch reference
+        with torch.no_grad():
+            pt_mul, pt_add, pt_enhanced = model(dummy)
+
+        # ONNX Runtime
+        ort_inputs = {"input": dummy.numpy()}
+        ort_mul, ort_add, ort_enhanced = session.run(None, ort_inputs)
+
+        # Compare enhanced output (the one that matters most)
+        for name, pt_out, ort_out in [
+            ("mul", pt_mul, ort_mul),
+            ("add", pt_add, ort_add),
+            ("enhanced", pt_enhanced, ort_enhanced),
+        ]:
+            max_diff = np.max(np.abs(pt_out.numpy() - ort_out))
+            if max_diff < 1e-5:
+                status = "OK"
+                symbol = "+"
+            elif max_diff < 1e-3:
+                status = "WARN"
+                symbol = "~"
+            else:
+                status = "FAIL"
+                symbol = "X"
+
+            print(f"  [{symbol}] {h}x{w} {name:10s} max_diff={max_diff:.2e} [{status}]")
+
+            if max_diff >= 1e-3:
+                print(f"      FAIL: max abs diff {max_diff:.6f} >= 1e-3")
+                all_ok = False
+
+    return all_ok
+
+
+# ===========================================================================
+# Main
+# ===========================================================================
+
+def main():
+    parser = argparse.ArgumentParser(description="Export IAT checkpoints to ONNX")
+    parser.add_argument(
+        "--checkpoints",
+        type=str,
+        default="all",
+        choices=["all", "exposure", "lol_v1", "lol_v2"],
+        help="Which checkpoint(s) to export",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default=str(Path(__file__).parent / "outputs"),
+        help="Directory for exported ONNX files",
+    )
+    parser.add_argument(
+        "--opset",
+        type=int,
+        default=17,
+        help="ONNX opset version",
+    )
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Determine which checkpoints to export
+    if args.checkpoints == "all":
+        names = list(CHECKPOINTS.keys())
+    else:
+        names = [args.checkpoints]
+
+    # Apply monkey-patches
+    print("=" * 60)
+    print("Applying ONNX-compatibility patches...")
+    print("=" * 60)
+    apply_patches()
+    print()
+
+    results = {}
+    for name in names:
+        cfg = CHECKPOINTS[name]
+        onnx_path = output_dir / f"iat_{name}.onnx"
+
+        print("=" * 60)
+        print(f"Exporting: {name} ({cfg['description']})")
+        print(f"  Checkpoint: {cfg['path']}")
+        print(f"  Output:     {onnx_path}")
+        print("=" * 60)
+
+        # Check checkpoint exists
+        if not cfg["path"].exists():
+            print(f"  SKIP: checkpoint not found at {cfg['path']}")
+            results[name] = "SKIP"
+            continue
+
+        # Load
+        t0 = time.time()
+        model = load_model(name)
+        print(f"  Loaded model in {time.time() - t0:.2f}s")
+
+        # Export
+        t0 = time.time()
+        export_onnx(model, onnx_path, args.opset)
+        export_time = time.time() - t0
+        file_size_mb = onnx_path.stat().st_size / (1024 * 1024)
+        print(f"  Exported in {export_time:.2f}s ({file_size_mb:.1f} MB)")
+
+        # Verify
+        print(f"  Verifying at {len(VERIFICATION_RESOLUTIONS)} resolutions...")
+        ok = verify_onnx(model, onnx_path)
+        results[name] = "PASS" if ok else "FAIL"
+        print()
+
+    # Summary
+    print("=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    all_pass = True
+    for name, status in results.items():
+        if status == "PASS":
+            symbol = "+"
+        elif status == "SKIP":
+            symbol = "-"
+        else:
+            symbol = "X"
+        print(f"  [{symbol}] {name}: {status}")
+        if status == "FAIL":
+            all_pass = False
+
+    if not all_pass:
+        print("\nSome exports FAILED numerical verification!")
+        sys.exit(1)
+    else:
+        print("\nAll exports passed!")
+        sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

onnx/iat_exposure.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd16c959a336b99ff15bef071ad6e9b081017241e5575d9e391653b676f2a5c6
+size 66820

onnx/iat_exposure.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f92c0ab92266a096d117d3747fbbce244ca91e8a3f28f70db5b2048f478418a8
+size 355264

onnx/iat_lol_v1.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfc50128b40b65edd0dbeeb724d283639d129f70fd838c159f9106c463277fca
+size 66698

onnx/iat_lol_v1.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc59e6ecbcff44683654c318b50e649321d83050528a14165f059b70d24d1901
+size 355264

onnx/iat_lol_v2.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:139d90e61bd8ae71dc7ac00c93a27c242f97876351f288713bc158cb5d04d3e5
+size 66698

onnx/iat_lol_v2.onnx.data

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fa9af18c4ae6e7f91ca41a895dec0188343bfcc4112e162ac92cee708f52cdcf
+size 355264

preprocessor_config.json

@@ -0,0 +1,5 @@
+{
+  "do_normalize": false,
+  "do_rescale": true,
+  "rescale_factor": 0.00392156862745098
+}