Upload folder using huggingface_hub

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Shilin Yan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,164 @@
----
-license: apache-2.0
----
+---
+license: mit
+library_name: pytorch
+tags:
+  - image-classification
+  - ai-generated-image-detection
+  - deepfake-detection
+  - computer-vision
+  - pytorch
+pipeline_tag: image-classification
+---
+
+# AIDE Image Detector
+
+This repository packages the `checkpoint-19.pth` model from the local AIDE training run at `/home/meet/Aivsre_001/AIDE/output_multisource_run1/checkpoint-19.pth` as a Hugging Face model repository. The exported weights are provided both as a safe deployment artifact in `model.safetensors` and, if uploaded, as the original PyTorch training snapshot `checkpoint-19.pth`.
+
+The model is based on **AIDE**: a hybrid AI-generated image detector that combines frequency-forensic evidence and high-level semantic cues. In this run, the detector uses:
+
+- A fixed **30-filter SRM high-pass bank** to expose subtle forensic residuals.
+- Two **ResNet-50-style frequency encoders** that process DCT-derived reconstructions.
+- A frozen **OpenCLIP ConvNeXt-XXL visual trunk** for high-level semantic/image-manifold features.
+- A final **MLP fusion head** that merges ConvNeXt and forensic embeddings into a binary classifier: `real` vs `fake`.
+
+## Architecture
+
+The forward path in [`models/AIDE.py`](./models/AIDE.py) is:
+
+1. Start from one RGB image.
+2. Build four DCT-based reconstructed views with [`data/dct.py`](./data/dct.py):
+   - `x_minmin`
+   - `x_maxmax`
+   - `x_minmin1`
+   - `x_maxmax1`
+3. Build a fifth view, `x_0`, from the normalized RGB image.
+4. Pass the four DCT views through the fixed SRM high-pass filters and into two ResNet branches.
+5. Pass the RGB view through the frozen OpenCLIP ConvNeXt-XXL trunk.
+6. Project the ConvNeXt pooled embedding from `3072 -> 256`.
+7. Average the four ResNet forensic embeddings into a single `2048`-dimensional frequency representation.
+8. Concatenate `[ConvNeXt_256, Forensic_2048]` into a `2304`-dimensional vector.
+9. Classify with an MLP `2304 -> 1024 -> 2`.
+
+Important implementation details taken directly from the code:
+
+- The frequency branch uses `HPF -> ResNet(Bottleneck, [3, 4, 6, 3])`.
+- The ConvNeXt branch is constructed with `open_clip.create_model_and_transforms("convnext_xxlarge", pretrained=None)` and then populated from the checkpoint weights.
+- The ConvNeXt trunk is frozen in the model definition used for this checkpoint.
+- Inside the model, the RGB input is remapped from ImageNet normalization to CLIP normalization before entering the ConvNeXt visual trunk.
+
+## Input Preparation
+
+Inference must follow the same preparation used during training/evaluation:
+
+1. Convert the image to RGB.
+2. Convert to tensor in `[0, 1]`.
+3. Use `DCT_base_Rec_Module(window_size=32, stride=16, output=256, grade_N=6)` to reconstruct four frequency-ranked views.
+4. Resize all five views to `256 x 256`.
+5. Normalize each view with:
+
+```python
+mean = [0.485, 0.456, 0.406]
+std = [0.229, 0.224, 0.225]
+```
+
+6. Stack the views in this exact order:
+
+```python
+[x_minmin, x_maxmax, x_minmin1, x_maxmax1, x_0]
+```
+
+The provided [`inference.py`](./inference.py) script reproduces this preparation pipeline.
+
+## Checkpoint Details
+
+- Source checkpoint: `checkpoint-19.pth`
+- Exported safe weights: `model.safetensors`
+- Labels:
+  - `0 -> real`
+  - `1 -> fake`
+- Epoch: `19`
+- Logged trainable parameters: `54,432,466`
+
+From the local training log for `output_multisource_run1`:
+
+- Epoch 19 validation/top-1 accuracy: `77.9186`
+- Epoch 19 validation loss: `0.4757`
+- Best validation/top-1 accuracy observed in the same run: `78.5831` at epoch `17`
+
+## Training Context
+
+This checkpoint came from a multi-source run configured with:
+
+- `data_path=/home/meet/Aivsre_001/aide_data/train_multi_v1`
+- `eval_data_path=/home/meet/Aivsre_001/aide_data/eval_multi_v1`
+- `epochs=20`
+- `batch_size=8`
+- `blr=5e-4`
+- `weight_decay=0.0`
+- `nb_classes=2`
+- `aa=rand-m9-mstd0.5-inc1`
+- `smoothing=0.1`
+
+The upstream AIDE project is introduced in the paper **"A Sanity Check for AI-generated Image Detection"** and uses a hybrid design intended to improve robustness on challenging real-world AI-image detection settings.
+
+## Files In This Repo
+
+- `model.safetensors`: exported model state dict for safer deployment.
+- `checkpoint-19.pth`: original PyTorch training snapshot, if uploaded.
+- `config.json`: architecture and label metadata.
+- `model.json`: lightweight manifest for this packaged repo.
+- `preprocessor_config.json`: image normalization and DCT-view preparation metadata.
+- `inference.py`: local loading and prediction helper.
+- `models/` and `data/`: source modules required to reconstruct the architecture.
+
+## Usage
+
+Clone or download the repository, then install dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Run local inference:
+
+```bash
+python inference.py --repo_dir . --image /path/to/image.jpg
+```
+
+Or use it programmatically:
+
+```python
+from PIL import Image
+
+from inference import load_model, predict_pil_images
+
+model = load_model(".")
+image = Image.open("example.jpg").convert("RGB")
+result = predict_pil_images(model, [image])[0]
+print(result)
+```
+
+Example output:
+
+```python
+{
+    "label": "fake",
+    "real_probability": 0.082134,
+    "fake_probability": 0.917866,
+}
+```
+
+## Notes
+
+- This repository is designed for **weight hosting and reproducible local inference**.
+- The architecture is custom and is not a native `transformers` `AutoModel` implementation.
+- Because the model relies on OpenCLIP ConvNeXt-XXL plus custom DCT/SRM preprocessing, users should use the provided loader and inference script.
+
+## Credits
+
+This packaged repository is derived from the original AIDE implementation:
+
+- Project: https://github.com/shilinyan99/AIDE
+- Paper: https://arxiv.org/abs/2406.19435
+
+Original AIDE code is MIT licensed. See [`LICENSE`](./LICENSE).

config.json

@@ -0,0 +1,60 @@
+{
+  "architectures": [
+    "AIDE"
+  ],
+  "model_type": "aide",
+  "library_name": "pytorch",
+  "task": "image-classification",
+  "num_labels": 2,
+  "id2label": {
+    "0": "real",
+    "1": "fake"
+  },
+  "label2id": {
+    "real": 0,
+    "fake": 1
+  },
+  "image_size": 256,
+  "input_tensor_shape": [
+    5,
+    3,
+    256,
+    256
+  ],
+  "frequency_branch": {
+    "hpf_filters": 30,
+    "resnet_backbones": 2,
+    "resnet_variant": "ResNet-50-style encoder",
+    "output_dim": 2048
+  },
+  "semantic_branch": {
+    "backbone": "OpenCLIP ConvNeXt-XXL visual trunk",
+    "pooled_dim": 3072,
+    "projection_dim": 256,
+    "frozen": true
+  },
+  "classifier": {
+    "type": "mlp",
+    "input_dim": 2304,
+    "hidden_dim": 1024,
+    "output_dim": 2
+  },
+  "dct_preparation": {
+    "window_size": 32,
+    "stride": 16,
+    "output_size": 256,
+    "grade_N": 6,
+    "selected_views": [
+      "x_minmin",
+      "x_maxmax",
+      "x_minmin1",
+      "x_maxmax1",
+      "x_0_rgb"
+    ]
+  },
+  "checkpoint": {
+    "source": "checkpoint-19.pth",
+    "epoch": 19
+  }
+}
+

data/__init__.py

@@ -0,0 +1,2 @@
+# Package marker for local DCT utilities.
+

data/dct.py

@@ -0,0 +1,149 @@
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+
+
+def DCT_mat(size):
+    m = [[ (np.sqrt(1./size) if i == 0 else np.sqrt(2./size)) * np.cos((j + 0.5) * np.pi * i / size) for j in range(size)] for i in range(size)]
+    return m
+
+def generate_filter(start, end, size):
+    return [[0. if i + j > end or i + j < start else 1. for j in range(size)] for i in range(size)]
+
+def norm_sigma(x):
+    return 2. * torch.sigmoid(x) - 1.
+
+class Filter(nn.Module):
+    def __init__(self, size, band_start, band_end, use_learnable=False, norm=False):
+        super(Filter, self).__init__()
+        self.use_learnable = use_learnable
+
+        self.base = nn.Parameter(torch.tensor(generate_filter(band_start, band_end, size)), requires_grad=False)
+        if self.use_learnable:
+            self.learnable = nn.Parameter(torch.randn(size, size), requires_grad=True)
+            self.learnable.data.normal_(0., 0.1)
+        self.norm = norm
+        if norm:
+            self.ft_num = nn.Parameter(torch.sum(torch.tensor(generate_filter(band_start, band_end, size))), requires_grad=False)
+
+
+    def forward(self, x):
+        if self.use_learnable:
+            filt = self.base + norm_sigma(self.learnable)
+        else:
+            filt = self.base
+
+        if self.norm:
+            y = x * filt / self.ft_num
+        else:
+            y = x * filt
+        return y
+
+class DCT_base_Rec_Module(nn.Module):
+    """_summary_
+
+    Args:
+        x: [C, H, W] -> [C*level, output, output]
+    """
+    def __init__(self, window_size=32, stride=16, output=256, grade_N=6, level_fliter=[0]):
+        super().__init__()
+        
+        assert output % window_size == 0
+        assert len(level_fliter) > 0
+        
+        self.window_size = window_size
+        self.grade_N = grade_N
+        self.level_N = len(level_fliter)
+        self.N = (output // window_size) * (output // window_size)
+        
+        self._DCT_patch = nn.Parameter(torch.tensor(DCT_mat(window_size)).float(), requires_grad=False)
+        self._DCT_patch_T = nn.Parameter(torch.transpose(torch.tensor(DCT_mat(window_size)).float(), 0, 1), requires_grad=False)
+        
+        self.unfold = nn.Unfold(
+            kernel_size=(window_size, window_size), stride=stride
+        )
+        self.fold0 = nn.Fold(
+            output_size=(window_size, window_size), 
+            kernel_size=(window_size, window_size), 
+            stride=window_size
+        )
+        
+        lm, mh = 2.82, 2
+        level_f = [
+            Filter(window_size, 0, window_size * 2)
+        ]
+        
+        self.level_filters = nn.ModuleList([level_f[i] for i in level_fliter])
+        self.grade_filters = nn.ModuleList([Filter(window_size, window_size * 2. / grade_N * i, window_size * 2. / grade_N * (i+1), norm=True) for i in range(grade_N)])
+        
+        
+    def forward(self, x):
+        
+        N = self.N
+        grade_N = self.grade_N
+        level_N = self.level_N
+        window_size = self.window_size
+        C, W, H = x.shape
+        x_unfold = self.unfold(x.unsqueeze(0)).squeeze(0)  
+        
+
+        _, L = x_unfold.shape
+        x_unfold = x_unfold.transpose(0, 1).reshape(L, C, window_size, window_size) 
+        x_dct = self._DCT_patch @ x_unfold @ self._DCT_patch_T
+        
+        y_list = []
+        for i in range(self.level_N):
+            x_pass = self.level_filters[i](x_dct)
+            y = self._DCT_patch_T @ x_pass @ self._DCT_patch
+            y_list.append(y)
+        level_x_unfold = torch.cat(y_list, dim=1)
+        
+        grade = torch.zeros(L).to(x.device)
+        w, k = 1, 2
+        for _ in range(grade_N):
+            _x = torch.abs(x_dct)
+            _x = torch.log(_x + 1)
+            _x = self.grade_filters[_](_x)
+            _x = torch.sum(_x, dim=[1,2,3])
+            grade += w * _x            
+            w *= k
+        
+        _, idx = torch.sort(grade)
+        max_idx = torch.flip(idx, dims=[0])[:N]
+        maxmax_idx = max_idx[0]
+        if len(max_idx) == 1:
+            maxmax_idx1 = max_idx[0]
+        else:
+            maxmax_idx1 = max_idx[1]
+
+        min_idx = idx[:N]
+        minmin_idx = idx[0]
+        if len(min_idx) == 1:
+            minmin_idx1 = idx[0]
+        else:
+            minmin_idx1 = idx[1]
+
+        x_minmin = torch.index_select(level_x_unfold, 0, minmin_idx)
+        x_maxmax = torch.index_select(level_x_unfold, 0, maxmax_idx)
+        x_minmin1 = torch.index_select(level_x_unfold, 0, minmin_idx1)
+        x_maxmax1 = torch.index_select(level_x_unfold, 0, maxmax_idx1)
+
+        x_minmin = x_minmin.reshape(1, level_N*C*window_size* window_size).transpose(0, 1)
+        x_maxmax = x_maxmax.reshape(1, level_N*C*window_size* window_size).transpose(0, 1)
+        x_minmin1 = x_minmin1.reshape(1, level_N*C*window_size* window_size).transpose(0, 1)
+        x_maxmax1 = x_maxmax1.reshape(1, level_N*C*window_size* window_size).transpose(0, 1)
+
+        x_minmin = self.fold0(x_minmin)
+        x_maxmax = self.fold0(x_maxmax)
+        x_minmin1 = self.fold0(x_minmin1)
+        x_maxmax1 = self.fold0(x_maxmax1)
+
+       
+        return x_minmin, x_maxmax, x_minmin1, x_maxmax1
+
+
+        
+
+

inference.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Iterable, List
+
+import numpy as np
+import torch
+from PIL import Image
+from safetensors.torch import load_file as load_safetensors
+from torchvision import transforms
+
+from data.dct import DCT_base_Rec_Module
+from models import AIDE as build_aide_model
+
+
+IMAGE_SIZE = 256
+TO_TENSOR = transforms.ToTensor()
+NORMALIZE_AND_RESIZE = transforms.Compose(
+    [
+        transforms.Resize([IMAGE_SIZE, IMAGE_SIZE]),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        ),
+    ]
+)
+
+
+def build_aide_input_from_pil(image: Image.Image, dct_module: DCT_base_Rec_Module) -> torch.Tensor:
+    image = image.convert("RGB")
+    image_tensor = TO_TENSOR(image)
+    x_minmin, x_maxmax, x_minmin1, x_maxmax1 = dct_module(image_tensor)
+
+    x_0 = NORMALIZE_AND_RESIZE(image_tensor)
+    x_minmin = NORMALIZE_AND_RESIZE(x_minmin)
+    x_maxmax = NORMALIZE_AND_RESIZE(x_maxmax)
+    x_minmin1 = NORMALIZE_AND_RESIZE(x_minmin1)
+    x_maxmax1 = NORMALIZE_AND_RESIZE(x_maxmax1)
+
+    return torch.stack([x_minmin, x_maxmax, x_minmin1, x_maxmax1, x_0], dim=0)
+
+
+def load_model(
+    repo_dir: str | Path,
+    device: str | None = None,
+    weights_name: str = "model.safetensors",
+) -> torch.nn.Module:
+    repo_dir = Path(repo_dir)
+    weights_path = repo_dir / weights_name
+    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = build_aide_model(resnet_path=None, convnext_path=None)
+    state_dict = load_safetensors(str(weights_path))
+    model.load_state_dict(state_dict, strict=True)
+    model.to(device)
+    model.eval()
+    return model
+
+
+@torch.inference_mode()
+def predict_pil_images(
+    model: torch.nn.Module,
+    images: Iterable[Image.Image],
+    device: str | None = None,
+) -> List[dict]:
+    device = device or next(model.parameters()).device.type
+    dct_module = DCT_base_Rec_Module()
+    batch = torch.stack([build_aide_input_from_pil(img, dct_module) for img in images], dim=0).to(device)
+    logits = model(batch)
+    probs = torch.softmax(logits, dim=-1).cpu().numpy()
+
+    outputs = []
+    for prob in probs:
+        real_prob = float(prob[0])
+        fake_prob = float(prob[1])
+        label = "fake" if fake_prob >= real_prob else "real"
+        outputs.append(
+            {
+                "label": label,
+                "real_probability": round(real_prob, 6),
+                "fake_probability": round(fake_prob, 6),
+            }
+        )
+    return outputs
+
+
+def _load_images(paths: Iterable[str]) -> List[Image.Image]:
+    return [Image.open(path).convert("RGB") for path in paths]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Run AIDE image detector inference.")
+    parser.add_argument("--repo_dir", type=str, default=".", help="Local path to the model repository.")
+    parser.add_argument("--image", type=str, nargs="+", required=True, help="One or more image paths.")
+    parser.add_argument("--device", type=str, default=None, help="cuda or cpu")
+    args = parser.parse_args()
+
+    model = load_model(args.repo_dir, device=args.device)
+    images = _load_images(args.image)
+    predictions = predict_pil_images(model, images, device=args.device)
+
+    for image_path, prediction in zip(args.image, predictions):
+        print(
+            {
+                "image": str(image_path),
+                **prediction,
+            }
+        )
+
+
+if __name__ == "__main__":
+    main()
+

model.json

@@ -0,0 +1,18 @@
+{
+  "name": "AIDE Image Detector",
+  "repo_id": "meet4150/AIDE_image_detector",
+  "framework": "PyTorch",
+  "weights": {
+    "primary": "model.safetensors",
+    "training_snapshot": "checkpoint-19.pth"
+  },
+  "config_file": "config.json",
+  "preprocessor_file": "preprocessor_config.json",
+  "inference_file": "inference.py",
+  "labels": [
+    "real",
+    "fake"
+  ],
+  "notes": "Hybrid AIDE checkpoint exported from /home/meet/Aivsre_001/AIDE/output_multisource_run1/checkpoint-19.pth."
+}
+

model.safetensors

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

models/AIDE.py

@@ -0,0 +1,297 @@
+import torch
+import torch.nn as nn
+import open_clip
+import numpy as np
+
+from .srm_filter_kernel import all_normalized_hpf_list
+
+class HPF(nn.Module):
+  def __init__(self):
+    super(HPF, self).__init__()
+
+    #Load 30 SRM Filters
+    all_hpf_list_5x5 = []
+
+    for hpf_item in all_normalized_hpf_list:
+      if hpf_item.shape[0] == 3:
+        hpf_item = np.pad(hpf_item, pad_width=((1, 1), (1, 1)), mode='constant')
+
+      all_hpf_list_5x5.append(hpf_item)
+
+    hpf_weight = torch.Tensor(all_hpf_list_5x5).view(30, 1, 5, 5).contiguous()
+    hpf_weight = torch.nn.Parameter(hpf_weight.repeat(1, 3, 1, 1), requires_grad=False)
+   
+
+    self.hpf = nn.Conv2d(3, 30, kernel_size=5, padding=2, bias=False)
+    self.hpf.weight = hpf_weight
+
+
+  def forward(self, input):
+
+    output = self.hpf(input)
+
+    return output
+
+
+
+def conv3x3(in_planes, out_planes, stride=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=1, bias=False)
+
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(BasicBlock, self).__init__()
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None):
+        super(Bottleneck, self).__init__()
+        self.conv1 = conv1x1(inplanes, planes)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.conv2 = conv3x3(planes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.conv3 = conv1x1(planes, planes * self.expansion)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=True):
+        super(ResNet, self).__init__()
+
+        self.inplanes = 64
+        self.conv1 = nn.Conv2d(30, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+
+
+        return x
+
+class Mlp(nn.Module):
+    """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+    """
+
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+class AIDE_Model(nn.Module):
+
+    def __init__(self, resnet_path, convnext_path):
+        super(AIDE_Model, self).__init__()
+        self.hpf = HPF()
+        self.model_min = ResNet(Bottleneck, [3, 4, 6, 3])
+        self.model_max = ResNet(Bottleneck, [3, 4, 6, 3])
+       
+        if resnet_path is not None:
+            pretrained_dict = torch.load(resnet_path, map_location='cpu')
+        
+            model_min_dict = self.model_min.state_dict()
+            model_max_dict = self.model_max.state_dict()
+    
+            for k in pretrained_dict.keys():
+                if k in model_min_dict and pretrained_dict[k].size() == model_min_dict[k].size():
+                    model_min_dict[k] = pretrained_dict[k]
+                    model_max_dict[k] = pretrained_dict[k]
+                else:
+                    print(f"Skipping layer {k} because of size mismatch")
+        
+        self.fc = Mlp(2048 + 256 , 1024, 2)
+
+        print("build model with convnext_xxl")
+        self.openclip_convnext_xxl, _, _ = open_clip.create_model_and_transforms(
+            "convnext_xxlarge", pretrained=convnext_path
+        )
+
+        self.openclip_convnext_xxl = self.openclip_convnext_xxl.visual.trunk
+        self.openclip_convnext_xxl.head.global_pool = nn.Identity()
+        self.openclip_convnext_xxl.head.flatten = nn.Identity()
+
+        self.openclip_convnext_xxl.eval()
+        
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.convnext_proj = nn.Sequential(
+            nn.Linear(3072, 256),
+
+        )
+        for param in self.openclip_convnext_xxl.parameters():
+            param.requires_grad = False
+
+    
+
+    def forward(self, x):
+
+        b, t, c, h, w = x.shape
+
+        x_minmin = x[:, 0] #[b, c, h, w]
+        x_maxmax = x[:, 1]
+        x_minmin1 = x[:, 2]
+        x_maxmax1 = x[:, 3]
+        tokens = x[:, 4]
+
+        x_minmin = self.hpf(x_minmin)
+        x_maxmax = self.hpf(x_maxmax)
+        x_minmin1 = self.hpf(x_minmin1)
+        x_maxmax1 = self.hpf(x_maxmax1)
+
+        with torch.no_grad():
+            
+            clip_mean = torch.Tensor([0.48145466, 0.4578275, 0.40821073])
+            clip_mean = clip_mean.to(tokens, non_blocking=True).view(3, 1, 1)
+            clip_std = torch.Tensor([0.26862954, 0.26130258, 0.27577711])
+            clip_std = clip_std.to(tokens, non_blocking=True).view(3, 1, 1)
+            dinov2_mean = torch.Tensor([0.485, 0.456, 0.406]).to(tokens, non_blocking=True).view(3, 1, 1)
+            dinov2_std = torch.Tensor([0.229, 0.224, 0.225]).to(tokens, non_blocking=True).view(3, 1, 1)
+
+            local_convnext_image_feats = self.openclip_convnext_xxl(
+                tokens * (dinov2_std / clip_std) + (dinov2_mean - clip_mean) / clip_std
+            ) #[b, 3072, 8, 8]
+            assert local_convnext_image_feats.size()[1:] == (3072, 8, 8)
+            local_convnext_image_feats = self.avgpool(local_convnext_image_feats).view(tokens.size(0), -1)
+            x_0 = self.convnext_proj(local_convnext_image_feats)
+
+        x_min = self.model_min(x_minmin)
+        x_max = self.model_max(x_maxmax)
+        x_min1 = self.model_min(x_minmin1)
+        x_max1 = self.model_max(x_maxmax1)
+
+        x_1 = (x_min + x_max + x_min1 + x_max1) / 4
+
+        x = torch.cat([x_0, x_1], dim=1)
+
+        x = self.fc(x)
+
+        return x
+
+def AIDE(resnet_path, convnext_path):
+    model = AIDE_Model(resnet_path, convnext_path)
+    return model
+

models/__init__.py

@@ -0,0 +1,2 @@
+from .AIDE import AIDE
+

models/srm_filter_kernel.py

@@ -0,0 +1,220 @@
+
+import numpy as np
+
+filter_class_1 = [
+  np.array([
+    [1, 0, 0],
+    [0, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 1, 0],
+    [0, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 1],
+    [0, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [1, -1, 0],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 1],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 0],
+    [1, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 0],
+    [0, 1, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [0, -1, 0],
+    [0, 0, 1]
+  ], dtype=np.float32)
+]
+
+
+filter_class_2 = [
+  np.array([
+    [1, 0, 0],
+    [0, -2, 0],
+    [0, 0, 1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 1, 0],
+    [0, -2, 0],
+    [0, 1, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 1],
+    [0, -2, 0],
+    [1, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [1, -2, 1],
+    [0, 0, 0]
+  ], dtype=np.float32),
+]
+
+
+filter_class_3 = [
+  np.array([
+    [-1, 0, 0, 0, 0],
+    [0, 3, 0, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 0, 1, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, -1, 0, 0],
+    [0, 0, 3, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 1, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, -1],
+    [0, 0, 0, 3, 0],
+    [0, 0, -3, 0, 0],
+    [0, 1, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0],
+    [0, 1, -3, 3, -1],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 1, 0, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 0, 3, 0],
+    [0, 0, 0, 0, -1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 1, 0, 0],
+    [0, 0, -3, 0, 0],
+    [0, 0, 3, 0, 0],
+    [0, 0, -1, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 1, 0],
+    [0, 0, -3, 0, 0],
+    [0, 3, 0, 0, 0],
+    [-1, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0],
+    [-1, 3, -3, 1, 0],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32)
+]
+
+
+filter_edge_3x3 = [
+  np.array([
+    [-1, 2, -1],
+    [2, -4, 2],
+    [0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 2, -1],
+    [0, -4, 2],
+    [0, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0],
+    [2, -4, 2],
+    [-1, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [-1, 2, 0],
+    [2, -4, 0],
+    [-1, 2, 0]
+  ], dtype=np.float32),
+]
+
+filter_edge_5x5 = [
+  np.array([
+    [-1, 2, -2, 2, -1],
+    [2, -6, 8, -6, 2],
+    [-2, 8, -12, 8, -2],
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, -2, 2, -1],
+    [0, 0, 8, -6, 2],
+    [0, 0, -12, 8, -2],
+    [0, 0, 8, -6, 2],
+    [0, 0, -2, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [0, 0, 0, 0, 0],
+    [0, 0, 0, 0, 0],
+    [-2, 8, -12, 8, -2],
+    [2, -6, 8, -6, 2],
+    [-1, 2, -2, 2, -1]
+  ], dtype=np.float32),
+  np.array([
+    [-1, 2, -2, 0, 0],
+    [2, -6, 8, 0, 0],
+    [-2, 8, -12, 0, 0],
+    [2, -6, 8, 0, 0],
+    [-1, 2, -2, 0, 0]
+  ], dtype=np.float32),
+]
+
+square_3x3 = np.array([
+  [-1, 2, -1],
+  [2, -4, 2],
+  [-1, 2, -1]
+], dtype=np.float32)
+
+square_5x5 = np.array([
+  [-1, 2, -2, 2, -1],
+  [2, -6, 8, -6, 2],
+  [-2, 8, -12, 8, -2],
+  [2, -6, 8, -6, 2],
+  [-1, 2, -2, 2, -1]
+], dtype=np.float32)
+
+
+all_hpf_list = filter_class_1 + filter_class_2 + filter_class_3 + filter_edge_3x3 + filter_edge_5x5 + [square_3x3, square_5x5]
+
+hpf_3x3_list = filter_class_1 + filter_class_2 + filter_edge_3x3 + [square_3x3]
+hpf_5x5_list = filter_class_3 + filter_edge_5x5 + [square_5x5]
+
+normalized_filter_class_2 = [hpf / 2 for hpf in filter_class_2]
+normalized_filter_class_3 = [hpf / 3 for hpf in filter_class_3]
+normalized_filter_edge_3x3 = [hpf / 4 for hpf in filter_edge_3x3]
+normalized_square_3x3 = square_3x3 / 4
+normalized_filter_edge_5x5 = [hpf / 12 for hpf in filter_edge_5x5]
+normalized_square_5x5 = square_5x5 / 12
+
+all_normalized_hpf_list = filter_class_1 + normalized_filter_class_2 + normalized_filter_class_3 + \
+ normalized_filter_edge_3x3 + normalized_filter_edge_5x5 + [normalized_square_3x3, normalized_square_5x5]
+
+normalized_hpf_3x3_list = filter_class_1 + normalized_filter_class_2 + normalized_filter_edge_3x3 + [normalized_square_3x3]
+normalized_hpf_5x5_list = normalized_filter_class_3 + normalized_filter_edge_5x5 + [normalized_square_5x5]
+
+normalized_3x3_list = normalized_filter_edge_3x3 + [normalized_square_3x3]
+normalized_5x5_list = normalized_filter_edge_5x5 + [normalized_square_5x5]

preprocessor_config.json

@@ -0,0 +1,35 @@
+{
+  "do_convert_rgb": true,
+  "do_resize": true,
+  "size": {
+    "height": 256,
+    "width": 256
+  },
+  "resample": 3,
+  "do_rescale": true,
+  "rescale_factor": 0.00392156862745098,
+  "do_normalize": true,
+  "image_mean": [
+    0.485,
+    0.456,
+    0.406
+  ],
+  "image_std": [
+    0.229,
+    0.224,
+    0.225
+  ],
+  "extra_views": {
+    "type": "dct_reconstruction",
+    "window_size": 32,
+    "stride": 16,
+    "grade_N": 6,
+    "selected_views": [
+      "x_minmin",
+      "x_maxmax",
+      "x_minmin1",
+      "x_maxmax1"
+    ]
+  }
+}
+

requirements.txt

@@ -0,0 +1,7 @@
+torch>=2.0
+torchvision>=0.15
+numpy>=1.24
+Pillow>=9.5
+safetensors>=0.4
+open-clip-torch>=2.24.0
+