Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,7 @@ 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
+figs/Main.png filter=lfs diff=lfs merge=lfs -text
+fonts/Rainbow-Party-2.ttf filter=lfs diff=lfs merge=lfs -text
+scripts/data_process/ruod/008431.jpg filter=lfs diff=lfs merge=lfs -text
+scripts/evaluation/FasterRCNN_score-mmdet/configs/reppoints/reppoints.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 CVTEAM
+
+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

@@ -0,0 +1,252 @@
+# Envisioning Beyond the Few: Disentangled Semantics and Primitives for Few-Shot Atypical Layout-to-Image Generation
+
+**ICML 2026**
+
+**Authors:** Nan Bao, Yifan Zhao, Wenzhuang Wang, Jia Li
+
+![Main](figs/Main.png)
+
+## Environment Setup
+
+We use two separate environments:
+
+1. **Main environment** for core training and inference.
+
+    ```bash
+    conda create -n dsp python=3.10.20
+    conda activate dsp
+    pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 --index-url https://download.pytorch.org/whl/cu126
+    pip install datasets==4.8.5 pillow==12.2.0 accelerate==1.13.0 transformers==5.8.1 diffusers==0.38.0 safetensors==0.8.0rc0 tensorboard==2.20.0 opencv-python==4.13.0.92 einops==0.8.2 imagesize==2.0.0 peft==0.19.1 ttach==0.0.3 ftfy==6.3.1 albumentations==2.0.8
+    ```
+
+2. **Evaluation environment** for MMDetection/MMEngine compatibility. It is used for evaluation with MMDetection/MMEngine due to strict version constraints, and also supports YOLO-based evaluation.
+
+    ```bash
+    conda create -n dsp-eval python=3.10.20
+    conda activate dsp-eval
+    conda install mkl==2023.1.0 numpy==1.26.4
+    conda install pytorch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 pytorch-cuda=12.1 -c pytorch -c nvidia
+    pip install mmengine==0.10.7 tqdm==4.67.3 shapely==2.1.2 scipy==1.15.3 terminaltables==3.1.10 ultralytics==8.4.50 pycocotools==2.0.11 https://download.openmmlab.com/mmcv/dist/cu121/torch2.1.0/mmcv-2.1.0-cp310-cp310-manylinux1_x86_64.whl "numpy<2.0.0" "setuptools<70.0.0"
+    ```
+
+## Set Environment Variables
+
+Set the root path of this project:
+
+```bash
+export DSP_PROJECT_DIR=/path/to/DSP # replace with the actual path
+```
+
+It is recommended to add this line to `~/.bashrc` or `~/.zshrc` for persistence.
+
+## Pretrained Models Preparation
+
+1. We use several pretrained models as external dependencies. Please download them manually from the following sources:
+    - [stable-diffusion-v1-5](https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5)
+    - [clip-vit-large-patch14](https://huggingface.co/openai/clip-vit-large-patch14)
+    - [dinov2_vitl14_pretrain.pth](https://dl.fbaipublicfiles.com/dinov2/dinov2_vitl14/dinov2_vitl14_pretrain.pth)
+    - [ViT-B-16.pt](https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt)
+
+2. After downloading, organize the pretrained weights under `./pretrained` as follows:
+
+    ```bash
+    pretrained
+    ├── stable-diffusion-v1-5
+    │   └── ...
+    ├── clip-vit-large-patch14
+    │   └── ...
+    ├── dinov2_vitl14_pretrain.pth
+    └── ViT-B-16.pt
+    ```
+
+    You may either copy or symlink the files. We recommend using symbolic links:
+
+    ```bash
+    ln -s /path/to/stable-diffusion-v1-5 ./pretrained/stable-diffusion-v1-5
+    ln -s /path/to/clip-vit-large-patch14 ./pretrained/clip-vit-large-patch14
+    ln -s /path/to/dinov2_vitl14_pretrain.pth ./pretrained/dinov2_vitl14_pretrain.pth
+    ln -s /path/to/ViT-B-16.pt ./pretrained/ViT-B-16.pt
+    ```
+
+## Data Preparation
+
+1. We use several public datasets. Please download them manually from the following sources:
+
+    - [DIOR](https://gcheng-nwpu.github.io/#Datasets)
+    - [RUOD](https://github.com/xiaoDetection/RUOD)
+    - [ExDark](https://github.com/cs-chan/Exclusively-Dark-Image-Dataset/tree/master/Dataset)
+
+2. Unzip the downloaded datasets and organize the external dataset directories as follows:
+
+    ```bash
+    DIOR-VOC
+    ├── Annotations
+    │   ├── Horizontal_Bounding_Boxes
+    │   └── Oriented_Bounding_Boxes
+    └── VOC2007
+        ├── ImageSets
+        │   ├── Layout
+        │   ├── Main
+        │   └── Segmentation
+        └── JPEGImages
+    ```
+
+    ```bash
+    RUOD
+    ├── Environment_pic
+    │   ├── blur
+    │   ├── color
+    │   └── light
+    ├── Environmet_ANN
+    ├── RUOD_ANN
+    └── RUOD_pic
+        ├── test
+        └── train
+    ```
+
+    ```bash
+    ExDark
+    ├── annos
+    ├── imageclasslist.txt
+    └── images
+    ```
+
+3. Run data preprocessing scripts located in `./scripts/data_process`, after updating all hard-coded paths (e.g., `/path/to/DIOR_VOC`, `/path/to/RUOD`, `/path/to/ExDark`) in the scripts to match the local setup. Execute them in order.
+
+    The preprocessing outputs will be generated under `./data` with the following structure:
+
+    ```bash
+    data
+    ├── DIOR
+    │   ├── dior_emb.pt
+    │   ├── images -> /path/to/DIOR-VOC/VOC2007/JPEGImages
+    │   ├── metadatas
+    │   └── patches
+    ├── EXDARK
+    │   ├── exdark_emb.pt
+    │   ├── images
+    │   ├── metadatas
+    │   └── patches
+    └── RUOD
+        ├── images -> /path/to/RUOD/RUOD_pic
+        ├── metadatas
+        ├── patches
+        └── ruod_emb.pt
+    ```
+
+## Training and Inference
+
+We provide three example configurations in `./configs`: `dsp-dior.yaml`, `dsp-ruod.yaml`, and `dsp-exdark.yaml`.
+
+> **Argument Description:**
+> - **config:** configuration file for model and dataset setup.
+> - **metaseed:** seed generator identifier for deterministic sampling.
+> - **num_seed:** number of sampling seeds for few-shot evaluation.
+> - **k_shot:** number of samples per category in few-shot setting.
+> - **run_id:** identifier for different runs.
+> - **gpu_ids:** GPU device indices for execution.
+> - **iter:** number of bootstrap iterations for FID.
+
+### Base Phase Training
+
+```bash
+bash train_base.sh --config "dsp-dior"
+bash train_base.sh --config "dsp-ruod"
+bash train_base.sh --config "dsp-exdark"
+```
+
+### Novel Phase Training
+
+```bash
+bash train_novel.sh --config "dsp-dior" --metaseed "aaa" --num_seed 50 --k_shot "5" --run_id "1" --gpu_ids "0,1,2,3"
+bash train_novel.sh --config "dsp-ruod" --metaseed "aaa" --num_seed 50 --k_shot "5" --run_id "1" --gpu_ids "0,1,2,3"
+bash train_novel.sh --config "dsp-exdark" --metaseed "aaa" --num_seed 50 --k_shot "5" --run_id "1" --gpu_ids "0,1,2,3"
+```
+
+### Inference
+
+```bash
+bash infer.sh --config "dsp-dior" --metaseed "aaa" --num_seed 50 --k_shot "5" --run_id "1" --ckpt "100" --gpu_ids "0,1,2,3" --max_infer_size 50
+bash infer.sh --config "dsp-ruod" --metaseed "aaa" --num_seed 50 --k_shot "5" --run_id "1" --ckpt "100" --gpu_ids "0,1,2,3" --max_infer_size 50
+bash infer.sh --config "dsp-exdark" --metaseed "aaa" --num_seed 50 --k_shot "5" --run_id "1" --ckpt "100" --gpu_ids "0,1,2,3" --max_infer_size 50
+```
+
+## Evaluation
+
+### Preparation
+
+Download the YOLO and Faster R-CNN weights from [this link](https://drive.google.com/drive/folders/1FWN02KEuGPdEkXv38MmT8-D4-uQcAn4_?usp=sharing). Place them under `./pretrained`. The expected directory structure is as follows:
+
+```bash
+pretrained
+├── evaluation
+│   ├── mmdet
+│   │   ├── faster_rcnn_r50_fpn_1x-dior
+│   │   │   └── epoch_12.pth
+│   │   ├── faster_rcnn_r50_fpn_1x-exdark
+│   │   │   └── epoch_12.pth
+│   │   └── faster_rcnn_r50_fpn_1x-ruod
+│   │       └── epoch_12.pth
+│   └── yolo
+│       └── best.pt
+└── ... (pretrained models for training)
+```
+
+### YOLO (mAP / AP50 / AP75)
+
+> **Note:** In yolo-wrapper-dior.sh, the `--xml_folder` path should be set to the DIOR annotation directory (`/path/to/DIOR-VOC/Annotations/Horizontal_Bounding_Boxes`).
+
+```bash
+cd $DSP_PROJECT_DIR/scripts/evaluation/yoloscore-dior
+bash yolo-wrapper-dior.sh --config "dsp-dior" --metaseed "aaa" --num_seed 50 --ckpt "100" --k_shot "5" --run_id "1" --gpu_ids 0
+```
+
+### Faster R-CNN (mAP / AP50 / AP75)
+
+```bash
+cd $DSP_PROJECT_DIR/scripts/evaluation/FasterRCNN_score-mmdet
+bash test-wrapper-dior.sh --config "dsp-dior" --metaseed "aaa" --num_seed 50 --ckpt "100" --k_shot "5" --run_id "1" --gpu_ids 0
+bash test-wrapper-ruod.sh --config "dsp-ruod" --metaseed "aaa" --num_seed 50 --ckpt "100" --k_shot "5" --run_id "1" --gpu_ids 0
+bash test-wrapper-exdark.sh --config "dsp-exdark" --metaseed "aaa" --num_seed 50 --ckpt "100" --k_shot "5" --run_id "1" --gpu_ids 0
+```
+
+### Bootstrap FID
+
+```bash
+cd $DSP_PROJECT_DIR/scripts/evaluation/bootstrap_fid
+python boot_fid-dior.py --config dsp-dior -run_id 1 -num_seeds 50 --iter 50 --k_shot 5
+python boot_fid-ruod.py --config dsp-ruod -run_id 1 -num_seeds 50 --iter 50 --k_shot 5
+python boot_fid-exdark.py --config dsp-exdark -run_id 1 -num_seeds 50 --iter 50 --k_shot 5
+```
+
+Bootstrap FID results will be saved under `./metrics/BootstrapFID`.
+
+### Detection Metric Summarization
+
+```bash
+cd $DSP_PROJECT_DIR/scripts/evaluation/summarize
+bash summarize-wrapper.sh --config "dsp-dior" --k_shot "5" --run_id "1" --ckpt "100" --metaseed "aaa" --num_seed 50
+bash summarize-wrapper.sh --config "dsp-ruod" --k_shot "5" --run_id "1" --ckpt "100" --metaseed "aaa" --num_seed 50
+bash summarize-wrapper.sh --config "dsp-exdark" --k_shot "5" --run_id "1" --ckpt "100" --metaseed "aaa" --num_seed 50
+```
+
+Detection evaluation results (mAP / AP50 / AP75, YOLO and Faster R-CNN) will be summarized in `./metrics`.
+
+## Acknowledgement
+
+Our work is based on [stable diffusion](https://github.com/compvis/stable-diffusion), [diffusers](https://github.com/huggingface/diffusers), [CLIP](https://github.com/openai/CLIP), [DINOv2](https://github.com/facebookresearch/dinov2), [CC-Diff](https://github.com/AZZMM/CC-Diff), [MIGC](https://github.com/limuloo/MIGC), [GradCAM](https://github.com/linyq2117/CLIP-ES), and [kmeans_pytorch](https://github.com/subhadarship/kmeans_pytorch). Thanks for these great projects!
+
+## Citation
+
+If you find our work useful for your research, please cite the following paper.
+
+```bib
+@inproceedings{
+    bao2026envisioning,
+    title={Envisioning Beyond the Few: Disentangled Semantics and Primitives for Few-Shot Atypical Layout-to-Image Generation},
+    author={Bao, Nan and Zhao, Yifan and Wang, Wenzhuang and Li, Jia},
+    booktitle={Forty-third International Conference on Machine Learning},
+    year={2026},
+    url={https://openreview.net/forum?id=Jva4wVEySO}
+}
+```

claim1.py

@@ -0,0 +1,59 @@
+"""Claim 1 verification: GICDM out-of-sample generated-point scaling (Eq. 1, Algorithm 1).
+
+Test scenario (Figure 1 of paper): real samples drawn from a 60/40 mixture of two
+hyperspheres with radii (r1, r2); generated samples from a mixture with swapped
+radii and proportions. The two sets are disjoint, so all fidelity/coverage metrics
+should score 0 in the ideal case. Standard metrics fail in high dimension due to
+hubness; GICDM-corrected metrics should remain ~0.
+"""
+import numpy as np
+import sys, os
+sys.path.insert(0, os.path.dirname(__file__))
+from gicdm_core import (pairwise_sq_dists, icdm_scaling, gicdm,
+                        clipped_density, clipped_coverage, hubness_stats)
+
+
+def sample_mixture_spheres(d, n, r1, r2, p1):
+    n1 = int(round(p1 * n))
+    n2 = n - n1
+    def sphere(r, m):
+        v = np.random.normal(size=(m, d))
+        v /= np.linalg.norm(v, axis=1, keepdims=True)
+        return v * r
+    return np.vstack([sphere(r1, n1), sphere(r2, n2)])
+
+
+def run(d, n, seed=0):
+    np.random.seed(seed)
+    Xr = sample_mixture_spheres(d, n, 3.0, 5.0, 0.6)
+    # swapped radii & proportions
+    Xg = sample_mixture_spheres(d, n, 5.0, 3.0, 0.4)
+    K1, K2 = 5, 50  # K2 = 10*K1
+
+    # raw metric
+    cd_raw = clipped_density(Xr, Xg, k=5)
+    cc_raw = clipped_coverage(Xr, Xg, k=5)
+
+    # GICDM-corrected
+    D_final, keep, Drr_gicdm = gicdm(Xr, Xg, K1, K2)
+    dissim = (D_final, Drr_gicdm)  # (generated-to-real, real-to-real) GICDM dissimilarities
+    cd_gicdm = clipped_density(Xr, Xg, k=5, dissim=dissim, keep=keep)
+    cc_gicdm = clipped_coverage(Xr, Xg, k=5, dissim=dissim, keep=keep)
+
+    # Hubness in raw real-real space
+    h5_raw, A5_raw = hubness_stats(pairwise_sq_dists(Xr), k=5)
+
+    return dict(d=d, n=n, cd_raw=cd_raw, cc_raw=cc_raw,
+                cd_gicdm=cd_gicdm, cc_gicdm=cc_gicdm,
+                kept=float(keep.mean()), h5_raw=h5_raw, A5_raw=A5_raw)
+
+
+if __name__ == "__main__":
+    import json
+    out = []
+    for d in [10, 50, 100, 500, 1000]:
+        r = run(d, 1000, seed=42)
+        out.append(r)
+        print(f"d={d:5d}  CD raw={r['cd_raw']:.3f} GICDM={r['cd_gicdm']:.3f} | "
+              f"CC raw={r['cc_raw']:.3f} GICDM={r['cc_gicdm']:.3f} | kept={r['kept']:.2f} h5={r['h5_raw']:.2f}")
+    json.dump(out, open("results/claim1_hypersphere.json", "w"))

claim2.py

@@ -0,0 +1,73 @@
+"""Claim 2 verification: Proposition 5.1 & Corollary 5.2.
+
+Prop 5.1: p_hat_{mu,K}(x_i) = 1/(N V_d mu_i^d) * (1/K Sum_k k^{1/d})^d is a local
+density estimator (asymptotically p(x_i)).
+Cor 5.2: at ICDM convergence, mu_i ~ mu_bar for all i => p_hat equal everywhere
+(density uniformized).
+
+We verify numerically:
+ (a) p_hat_{mu,K} correlates with the true density in a mixture of Gaussians.
+ (b) after ICDM, the spread of mu_i collapses (std -> ~0) while raw std is large,
+     confirming density gradient removal.
+"""
+import numpy as np
+import sys, os
+sys.path.insert(0, os.path.dirname(__file__))
+from gicdm_core import pairwise_sq_dists, icdm_scaling, knn_distances
+
+
+def unit_ball_volume(d):
+    from math import gamma, pi
+    return (pi ** (d / 2)) / gamma(d / 2 + 1)
+
+
+def p_hat_mu_k(X, K):
+    D = pairwise_sq_dists(X)
+    mu = knn_distances(D, K, self_included=True).mean(axis=1)
+    N, d = X.shape
+    Vd = unit_ball_volume(d)
+    w = (np.sum([k ** (1.0 / d) for k in range(1, K + 1)]) / K) ** d
+    return 1.0 / (N * Vd * mu ** d) * w, mu
+
+
+def run():
+    rng = np.random.default_rng(0)
+    # Mixture of two Gaussians with different variances -> different densities
+    d = 20
+    n1, n2 = 800, 800
+    X = np.vstack([
+        rng.normal(0, 0.5, size=(n1, d)),
+        rng.normal(3, 2.0, size=(n2, d)),
+    ])
+    true_dens = np.concatenate([
+        (2 * np.pi * 0.25) ** (-d / 2) * np.ones(n1),
+        (2 * np.pi * 4.0) ** (-d / 2) * np.ones(n2),
+    ])
+    K = 20
+    phat, mu = p_hat_mu_k(X, K)
+    # correlation between estimator and true density (log scale, robust)
+    log_corr = np.corrcoef(np.log(phat), np.log(true_dens))[0, 1]
+
+    # ICDM convergence: spread of mu before/after
+    mu_before = mu.copy()
+    delta, mu_after = icdm_scaling(pairwise_sq_dists(X), K, n_iter=10, return_mu=True)
+    res = dict(
+        d=d, N=len(X), K=K,
+        log_corr_density=float(log_corr),
+        mu_before_std=float(mu_before.std()),
+        mu_before_cv=float(mu_before.std() / mu_before.mean()),
+        mu_after_std=float(mu_after.std()),
+        mu_after_cv=float(mu_after.std() / mu_after.mean()),
+    )
+    print("Prop 5.1 log-correlation p_hat vs true density:", round(res['log_corr_density'], 3))
+    print("mu_i std before ICDM:", round(res['mu_before_std'], 4),
+          "| after ICDM:", round(res['mu_after_std'], 5))
+    print("mu_i CV  before ICDM:", round(res['mu_before_cv'], 4),
+          "| after ICDM:", round(res['mu_after_cv'], 5))
+    return res
+
+
+if __name__ == "__main__":
+    import json
+    res = run()
+    json.dump(res, open("results/claim2_density.json", "w"))

claim5.py

@@ -0,0 +1,107 @@
+"""Claim 5 verification (scaled reproduction): Raisa et al. (2025) synthetic benchmark.
+
+We implement a subset of the Raisa et al. test scenarios that exercise the two
+metrics the paper reports gains on (Clipped Density & Clipped Coverage):
+  - GAUSSIAN MEAN DIFFERENCE (Purpose + Bounds: score should be 1 at equality)
+  - GAUSSIAN STD DEVIATION DIFFERENCE (Purpose: fails without GICDM)
+  - HYPERSPHERE SURFACE (Purpose: real vs generated on sphere shell)
+  - MODE COLLAPSE (Purpose / Bounds)
+  - SPHERE VS TORUS
+
+For each scenario we evaluate Clipped Density and Clipped Coverage with and without
+GICDM and check whether the metric behaves as the test expects. We tally pass rates
+across the implemented scenarios and compare the *direction* of the paper's reported
+gains (8/14->10/14 Purpose for Clipped Density; 8/13->11/13 Bounds).
+
+NOTE: this is a scaled reproduction of the statistical mechanism (full 14 Purpose /
+13 Bounds tests require the official benchmark suite). It validates that GICDM
+removes hubness-induced failures on representative scenarios.
+"""
+import numpy as np
+import sys, os
+sys.path.insert(0, os.path.dirname(__file__))
+from gicdm_core import pairwise_sq_dists, gicdm, clipped_density, clipped_coverage
+
+
+def run_gicdm_metrics(Xr, Xg, k=5, K1=5, K2=50):
+    Df, keep, Drr_g = gicdm(Xr, Xg, K1, K2)
+    cd = clipped_density(Xr, Xg, k=k, dissim=(Df, Drr_g), keep=keep)
+    cc = clipped_coverage(Xr, Xg, k=k, dissim=(Df, Drr_g), keep=keep)
+    cd0 = clipped_density(Xr, Xg, k=k)
+    cc0 = clipped_coverage(Xr, Xg, k=k)
+    return dict(cd0=cd0, cc0=cc0, cd=cd, cc=cc)
+
+
+def scenario_gauss_mean(d=100, n=1500, shift=0.0, seed=0):
+    rng = np.random.default_rng(seed)
+    Xr = rng.normal(0, 1, size=(n, d))
+    Xg = rng.normal(shift, 1, size=(n, d))
+    return Xr, Xg
+
+
+def scenario_gauss_std(d=100, n=1500, s=1.0, seed=0):
+    rng = np.random.default_rng(seed)
+    Xr = rng.normal(0, 1, size=(n, d))
+    Xg = rng.normal(0, s, size=(n, d))
+    return Xr, Xg
+
+
+def scenario_hypersphere(d=100, n=1500, r_r=1.0, r_g=1.0, seed=0):
+    rng = np.random.default_rng(seed)
+    def shell(r, m):
+        v = rng.normal(size=(m, d)); v /= np.linalg.norm(v, axis=1, keepdims=True)
+        return v * r
+    return shell(r_r, n), shell(r_g, n)
+
+
+def scenario_mode_collapse(d=100, n=1500, n_modes=5, collapse=False, seed=0):
+    rng = np.random.default_rng(seed)
+    centers = rng.normal(0, 5, size=(n_modes, d))
+    if collapse:
+        # generated collapses to a single mode
+        c = centers[0]
+        Xg = c[None, :] + rng.normal(0, 0.3, size=(n, d))
+    else:
+        sel = rng.integers(0, n_modes, size=n)
+        Xg = centers[sel] + rng.normal(0, 0.3, size=(n, d))
+    selr = rng.integers(0, n_modes, size=n)
+    Xr = centers[selr] + rng.normal(0, 0.3, size=(n, d))
+    return Xr, Xg
+
+
+def main():
+    results = {}
+    # GAUSSIAN MEAN DIFFERENCE: at shift=0 generated==real -> CD/CC should be ~1 (ideal)
+    Xr, Xg = scenario_gauss_mean(shift=0.0)
+    m = run_gicdm_metrics(Xr, Xg)
+    results['gauss_mean_ideal'] = m  # expect cd0,cd ~1 (pass)
+
+    # GAUSSIAN STD DEVIATION DIFFERENCE: mismatch in scale -> raw metric distorted in HD
+    Xr, Xg = scenario_gauss_std(s=1.5)
+    m = run_gicdm_metrics(Xr, Xg)
+    results['gauss_std'] = m
+
+    # HYPERSPHERE SURFACE: real on r=1, generated on r=1.0 -> equal -> CD ~1
+    Xr, Xg = scenario_hypersphere(r_r=1.0, r_g=1.0)
+    m = run_gicdm_metrics(Xr, Xg)
+    results['hypersphere_equal'] = m
+
+    # MODE COLLAPSE: generated collapses to 1 mode -> coverage should drop
+    Xr, Xg = scenario_mode_collapse(collapse=True)
+    m = run_gicdm_metrics(Xr, Xg)
+    results['mode_collapse'] = m
+
+    # MODE collapse absent -> coverage ~1
+    Xr, Xg = scenario_mode_collapse(collapse=False)
+    m = run_gicdm_metrics(Xr, Xg)
+    results['mode_full'] = m
+
+    for k_, v in results.items():
+        print(f"{k_:22s} CD raw={v['cd0']:.3f} GICDM={v['cd']:.3f} | "
+              f"CC raw={v['cc0']:.3f} GICDM={v['cc']:.3f}")
+    import json
+    json.dump(results, open("results/claim5_benchmark.json", "w"))
+
+
+if __name__ == "__main__":
+    main()

claim5_official.py

@@ -0,0 +1,79 @@
+"""Claim 5 (faithful, scaled): Raisa et al. synthetic benchmark mechanism.
+
+Uses the OFFICIAL GICDM repo (github.com/nicolassalvy/GICDM) Clipped Density /
+Clipped Coverage with the standard vs GICDM DataProcessor. We run representative
+Raisa scenarios on modest synthetic data (scaled down N,d) and tally whether each
+metric behaves correctly with/without GICDM. We report the *direction* of change
+and note the paper's full pass counts (Clipped Density Purpose 8/14 -> 10/14,
+Bounds 8/13 -> 11/13; Clipped Coverage 8/14->10/14, 9/13->11/13).
+"""
+import numpy as np
+import sys, os
+sys.path.insert(0, "/Users/equan_p/Developer/playground/ICML-2/official/GICDM")
+from metrics.hubness_processor.standard import DataProcessorStandard
+from metrics.hubness_processor.gicdm import GICDM
+from metrics.metrics.clipped_density_coverage import ClippedDensityCoverage
+
+
+def make_processor(Xr, use_gicdm, K=5):
+    if use_gicdm:
+        return GICDM(Xr, K=K, n_jobs=4, scale_factor=10)
+    return DataProcessorStandard(Xr, K=K, n_jobs=4)
+
+
+def cd_cc(Xr, Xg, use_gicdm, K=5):
+    dp = make_processor(Xr, use_gicdm, K)
+    cdc = ClippedDensityCoverage(dp)
+    cd = cdc.clipped_density(Xg)
+    cc = cdc.clipped_coverage(Xg)
+    return float(cd), float(cc)
+
+
+def sphere(d, n, r, rng):
+    v = rng.normal(size=(n, d)); v /= np.linalg.norm(v, axis=1, keepdims=True)
+    return v * r
+
+
+def main():
+    rng = np.random.default_rng(0)
+    d, n = 100, 1500
+    rows = []
+    scenarios = {}
+
+    # 1. GAUSSIAN MEAN DIFFERENCE: shift=0 -> identical -> CD,CC should be ~1
+    Xr = rng.normal(0, 1, size=(n, d)); Xg0 = rng.normal(0, 1, size=(n, d))
+    scenarios['gauss_mean_equal'] = (Xr.copy(), Xg0.copy())
+
+    # 2. GAUSSIAN STD DEVIATION DIFFERENCE: scale mismatch
+    Xg_std = rng.normal(0, 1.5, size=(n, d))
+    scenarios['gauss_std_1p5'] = (Xr.copy(), Xg_std.copy())
+
+    # 3. HYPERSPHERE SURFACE equal radius (should be ~1)
+    sr = sphere(d, n, 1.0, rng); sg = sphere(d, n, 1.0, rng)
+    scenarios['hypersphere_equal'] = (sr, sg)
+
+    # 4. MODE COLLAPSE: gen collapses to one of 5 modes
+    centers = rng.normal(0, 5, size=(5, d))
+    Xr_m = centers[rng.integers(0, 5, n)] + rng.normal(0, 0.3, (n, d))
+    Xg_mc = centers[0] + rng.normal(0, 0.3, (n, d))
+    scenarios['mode_collapse'] = (Xr_m.copy(), Xg_mc.copy())
+    Xg_mf = centers[rng.integers(0, 5, n)] + rng.normal(0, 0.3, (n, d))
+    scenarios['mode_full'] = (Xr_m.copy(), Xg_mf.copy())
+
+    # 5. SPHERE vs TORUS-like (sphere vs scaled sphere offset)
+    st = sphere(d, n, 1.0, rng) + 3.0  # offset sphere -> out of manifold
+    scenarios['sphere_offset'] = (sr.copy(), st.copy())
+
+    results = {}
+    for name, (Xr_, Xg_) in scenarios.items():
+        cd0, cc0 = cd_cc(Xr_, Xg_, False)
+        cdg, ccg = cd_cc(Xr_, Xg_, True)
+        results[name] = dict(cd_raw=cd0, cc_raw=cc0, cd_gicdm=cdg, cc_gicdm=ccg)
+        print(f"{name:18s} CD raw={cd0:.3f} GICDM={cdg:.3f} | CC raw={cc0:.3f} GICDM={ccg:.3f}")
+
+    import json
+    json.dump(results, open("/Users/equan_p/Developer/playground/ICML-2/repro_gicdm/results/claim5_official.json", "w"))
+
+
+if __name__ == "__main__":
+    main()

claim6.py

@@ -0,0 +1,34 @@
+"""Claim 6 verification: hubness statistics h5_1(1%) and A5 (Table 4).
+
+Reproduce Figure 7's Gaussian hubness-evolution: as dimension d increases, h5_1(1%)
+and A5 rise (hubness appears). Confirm the landmark reference values:
+ - no-hubness baseline (low d Gaussian): h5_1(1%) slightly above 2, A5 < 0.01
+ - high-d Gaussian: h5_1(1%) several, A5 ~0.1+  (matches ImageNet/DINOv2 order)
+
+We also reproduce Table 4's cross-embedding ordering qualitatively by noting the
+paper's empirical numbers; full reproduction of the exact 16 embeddings requires the
+official datasets/checkpoints (see logbook: github.com/nicolassalvy/GICDM). Here we
+validate the *statistical methodology* and the baseline/no-hubness reference.
+"""
+import numpy as np
+import sys, os
+sys.path.insert(0, os.path.dirname(__file__))
+from gicdm_core import pairwise_sq_dists, hubness_stats
+
+
+def gaussian_hubness(N=20000, dims=(10, 20, 50, 100, 200, 500, 1000, 2000), seed=0):
+    rng = np.random.default_rng(seed)
+    rows = []
+    for d in dims:
+        X = rng.normal(0, 1, size=(N, d))
+        D = pairwise_sq_dists(X)
+        h5, A5 = hubness_stats(D, k=5)
+        rows.append(dict(d=int(d), N=N, h5=float(h5), A5=float(A5)))
+        print(f"d={d:5d}  h5_1(1%)={h5:.2f}  A5={A5:.3f}")
+    return rows
+
+
+if __name__ == "__main__":
+    import json
+    rows = gaussian_hubness()
+    json.dump(rows, open("results/claim6_gaussian_hubness.json", "w"))

configs/dsp-dior.yaml

@@ -0,0 +1,99 @@
+seed: 42
+task_name: dsp-dior
+ckpt_dir: ./ckpt
+accelerator:
+  gradient_accumulation_steps: 80
+  mixed_precision: 'no'
+  report_to: tensorboard
+model:
+  name: dsp
+  sd15_weight_path: ./pretrained/stable-diffusion-v1-5
+  clip_weight_path: ./pretrained/clip-vit-large-patch14
+  dinov2_vitl14_path: ./pretrained/dinov2_vitl14_pretrain.pth
+  clip_vit_b16_path: ./pretrained/ViT-B-16.pt
+  exemplar_pool:
+    data_embeds_dict_path: ./data/DIOR/dior_emb.pt
+    exemplar_pool_path: ./data/DIOR/images
+  image_proj_model:
+    dim: 1280
+    depth: 4
+    dim_head: 64
+    num_queries: [16, 8, 8]
+    ff_mult: 4
+dataset:
+  name: dior
+  config: default
+  resolution: 512
+  ref_resolution: 224
+  categories:
+    all: [
+      vehicle, baseballfield, groundtrackfield, windmill, bridge, 
+      overpass, ship, airplane, tenniscourt, airport, 
+      expressway-service-area, basketballcourt, stadium, storagetank, chimney, 
+      dam, expressway-toll-station, golffield, trainstation, harbor
+    ]
+    base: [
+      vehicle, baseballfield, groundtrackfield, bridge, overpass, 
+      ship, airplane, tenniscourt, expressway-service-area, basketballcourt, 
+      stadium, storagetank, expressway-toll-station, golffield, harbor
+    ]
+    novel: [windmill, airport, chimney, dam, trainstation]
+  data_files:
+    train:
+      base:
+        default: ./data/DIOR/metadatas/data_setting1/train_base.jsonl
+      novel:
+        airport: ./data/DIOR/metadatas/data_setting1/train_novel_airport.jsonl
+        chimney: ./data/DIOR/metadatas/data_setting1/train_novel_chimney.jsonl
+        dam: ./data/DIOR/metadatas/data_setting1/train_novel_dam.jsonl
+        trainstation: ./data/DIOR/metadatas/data_setting1/train_novel_trainstation.jsonl
+        windmill: ./data/DIOR/metadatas/data_setting1/train_novel_windmill.jsonl
+    infer:
+      base:
+        default: ./data/DIOR/metadatas/data_setting1/test_base.jsonl
+      novel:
+        airport: ./data/DIOR/metadatas/data_setting1/test_novel_airport.jsonl
+        chimney: ./data/DIOR/metadatas/data_setting1/test_novel_chimney.jsonl
+        dam: ./data/DIOR/metadatas/data_setting1/test_novel_dam.jsonl
+        trainstation: ./data/DIOR/metadatas/data_setting1/test_novel_trainstation.jsonl
+        windmill: ./data/DIOR/metadatas/data_setting1/test_novel_windmill.jsonl
+  column_names: [image, captions, bndboxes, obboxes]
+  novel_settings:
+    k_shot: 5
+    shuffle_seed: 42
+  image_patch_path: ./data/DIOR/patches
+  transform: LayoutTransform
+training:
+  entry: train_dsp
+  optimizer:
+    name: AdamWGating
+    learning_rate: 1.e-4
+    weight_decay: 1.e-2
+    adam_beta: [0.9, 0.999]
+    adam_epsilon: 1.e-08
+  scheduler:
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+  batch_size: 1
+  num_workers: 0
+  noise_offset: 0
+  input_perturbation: 0
+  prediction_type: null
+  max_grad_norm: 1.
+  base:
+    num_train_epochs: 100
+    max_train_steps: null
+    ckpt_interval_steps: 1200
+  novel:
+    num_train_epochs: null
+    max_train_steps: 100
+    ckpt_interval_steps: 100
+    base_ckpt_steps: 1200
+inference:
+  entry: infer_dsp
+  output_dir: ./outputs
+  seed: 42
+  base:
+    ckpt_steps: 1200
+  novel:
+    ckpt_steps: 100

configs/dsp-exdark.yaml

@@ -0,0 +1,91 @@
+seed: 42
+task_name: dsp-exdark
+ckpt_dir: ./ckpt
+accelerator:
+  gradient_accumulation_steps: 80
+  mixed_precision: 'no'
+  report_to: tensorboard
+model:
+  name: dsp
+  sd15_weight_path: ./pretrained/stable-diffusion-v1-5
+  clip_weight_path: ./pretrained/clip-vit-large-patch14
+  dinov2_vitl14_path: ./pretrained/dinov2_vitl14_pretrain.pth
+  clip_vit_b16_path: ./pretrained/ViT-B-16.pt
+  exemplar_pool:
+    data_embeds_dict_path: ./data/EXDARK/exdark_emb.pt
+    exemplar_pool_path: ./data/EXDARK/images/train
+  image_proj_model:
+    dim: 1280
+    depth: 4
+    dim_head: 64
+    num_queries: [16, 8, 8]
+    ff_mult: 4
+dataset:
+  name: dior
+  config: default
+  resolution: 512
+  ref_resolution: 224
+  categories:
+    all: [
+      bicycle, boat, bottle, bus, car, cat,
+      chair, cup, dog, motorbike, people, table
+    ]
+    base: [bicycle, boat, bottle, car, cat, chair, cup, people]
+    novel: [bus, dog, motorbike, table]
+  data_files:
+    train:
+      base:
+        default: ./data/EXDARK/metadatas/data_setting1/train_base.jsonl
+      novel:
+        bus: ./data/EXDARK/metadatas/data_setting1/train_novel_bus.jsonl
+        dog: ./data/EXDARK/metadatas/data_setting1/train_novel_dog.jsonl
+        motorbike: ./data/EXDARK/metadatas/data_setting1/train_novel_motorbike.jsonl
+        table: ./data/EXDARK/metadatas/data_setting1/train_novel_table.jsonl
+    infer:
+      base:
+        default: [./data/EXDARK/metadatas/data_setting1/test_base.jsonl, ./data/EXDARK/metadatas/data_setting1/val_base.jsonl]
+      novel:
+        bus: [./data/EXDARK/metadatas/data_setting1/val_novel_bus.jsonl, ./data/EXDARK/metadatas/data_setting1/test_novel_bus.jsonl]
+        dog: [./data/EXDARK/metadatas/data_setting1/val_novel_dog.jsonl, ./data/EXDARK/metadatas/data_setting1/test_novel_dog.jsonl]
+        motorbike: [./data/EXDARK/metadatas/data_setting1/val_novel_motorbike.jsonl, ./data/EXDARK/metadatas/data_setting1/test_novel_motorbike.jsonl]
+        table: [./data/EXDARK/metadatas/data_setting1/val_novel_table.jsonl, ./data/EXDARK/metadatas/data_setting1/test_novel_table.jsonl]
+  column_names: [image, captions, bndboxes, obboxes]
+  novel_settings:
+    k_shot: 5
+    shuffle_seed: 42
+  image_patch_path: ./data/EXDARK/patches
+  transform: LayoutTransform
+training:
+  entry: train_dsp
+  optimizer:
+    name: AdamWGating
+    learning_rate: 1.e-4
+    weight_decay: 1.e-2
+    adam_beta: [0.9, 0.999]
+    adam_epsilon: 1.e-08
+  scheduler:
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+  batch_size: 1
+  num_workers: 0
+  noise_offset: 0
+  input_perturbation: 0
+  prediction_type: null
+  max_grad_norm: 1.
+  base:
+    num_train_epochs: 100
+    max_train_steps: null
+    ckpt_interval_steps: 600
+  novel:
+    num_train_epochs: null
+    max_train_steps: 100
+    ckpt_interval_steps: 100
+    base_ckpt_steps: 600
+inference:
+  entry: infer_dsp
+  output_dir: ./outputs
+  seed: 42
+  base:
+    ckpt_steps: 600
+  novel:
+    ckpt_steps: 100

configs/dsp-ruod.yaml

@@ -0,0 +1,91 @@
+seed: 42
+task_name: dsp-ruod
+ckpt_dir: ./ckpt
+accelerator:
+  gradient_accumulation_steps: 80
+  mixed_precision: 'no'
+  report_to: tensorboard
+model:
+  name: dsp
+  sd15_weight_path: ./pretrained/stable-diffusion-v1-5
+  clip_weight_path: ./pretrained/clip-vit-large-patch14
+  dinov2_vitl14_path: ./pretrained/dinov2_vitl14_pretrain.pth
+  clip_vit_b16_path: ./pretrained/ViT-B-16.pt
+  exemplar_pool:
+    data_embeds_dict_path: ./data/RUOD/ruod_emb.pt
+    exemplar_pool_path: ./data/RUOD/images/train
+  image_proj_model:
+    dim: 1280
+    depth: 4
+    dim_head: 64
+    num_queries: [16, 8, 8]
+    ff_mult: 4
+dataset:
+  name: dior
+  config: default
+  resolution: 512
+  ref_resolution: 224
+  categories:
+    all: [
+      holothurian, echinus, scallop, starfish, fish, 
+      corals, diver, cuttlefish, turtle, jellyfish, 
+    ]
+    base: [holothurian, echinus, scallop, starfish, fish, diver]
+    novel: [corals, cuttlefish, turtle, jellyfish]
+  data_files:
+    train:
+      base:
+        default: ./data/RUOD/metadatas/data_setting1/train_base.jsonl
+      novel:
+        corals: ./data/RUOD/metadatas/data_setting1/train_novel_corals.jsonl
+        cuttlefish: ./data/RUOD/metadatas/data_setting1/train_novel_cuttlefish.jsonl
+        turtle: ./data/RUOD/metadatas/data_setting1/train_novel_turtle.jsonl
+        jellyfish: ./data/RUOD/metadatas/data_setting1/train_novel_jellyfish.jsonl
+    infer:
+      base:
+        default: ./data/RUOD/metadatas/data_setting1/test_base.jsonl
+      novel:
+        corals: ./data/RUOD/metadatas/data_setting1/test_novel_corals.jsonl
+        cuttlefish: ./data/RUOD/metadatas/data_setting1/test_novel_cuttlefish.jsonl
+        turtle: ./data/RUOD/metadatas/data_setting1/test_novel_turtle.jsonl
+        jellyfish: ./data/RUOD/metadatas/data_setting1/test_novel_jellyfish.jsonl
+  column_names: [image, captions, bndboxes, obboxes]
+  novel_settings:
+    k_shot: 5
+    shuffle_seed: 42
+  image_patch_path: ./data/RUOD/patches
+  transform: LayoutTransform
+training:
+  entry: train_dsp
+  optimizer:
+    name: AdamWGating
+    learning_rate: 1.e-4
+    weight_decay: 1.e-2
+    adam_beta: [0.9, 0.999]
+    adam_epsilon: 1.e-08
+  scheduler:
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+  batch_size: 1
+  num_workers: 0
+  noise_offset: 0
+  input_perturbation: 0
+  prediction_type: null
+  max_grad_norm: 1.
+  base:
+    num_train_epochs: 100
+    max_train_steps: null
+    ckpt_interval_steps: 1200
+  novel:
+    num_train_epochs: null
+    max_train_steps: 100
+    ckpt_interval_steps: 100
+    base_ckpt_steps: 1200
+inference:
+  entry: infer_dsp
+  output_dir: ./outputs
+  seed: 42
+  base:
+    ckpt_steps: 1200
+  novel:
+    ckpt_steps: 100

crossover.py

@@ -0,0 +1,42 @@
+"""Proposition 5.3: crossover dimension d* for standard Gaussian N(0, I_d).
+
+d* solves  P( median NND_k^2 <= E||X||^2 ) = 1/2, i.e. the integral
+  int_0^infty P( Bin(N-1, F_{chi2_d}(lambda=r)) >= k ) f_{chi2_d}(r) dr = 1/2
+where F_{chi2_d}(lambda) is the noncentral chi2 CDF with noncentrality lambda,
+f_{chi2_d} is the chi2 density (central, lambda=0), and r = ||X||^2 ~ chi2_d.
+"""
+import numpy as np
+from scipy import integrate, stats
+
+
+def crossover_dimension(N, k, d_grid=None):
+    """Numerically solve Proposition 5.3 for d* given N and k."""
+    if d_grid is None:
+        d_grid = np.arange(2, 4001)  # up to 4000 dims
+
+    def median_prob(d):
+        # P( Bin(N-1, F_{chi2_d}(lambda=r)) >= k ) averaged over r ~ chi2_d
+        def integrand(r):
+            # noncentral chi2 cdf with noncentrality lambda = r
+            # F_{chi2_d}(lambda=r)(d) = P( chi2_d(r) <= d )
+            p = stats.ncx2.cdf(d, d, r)
+            return stats.binom.cdf(k - 1, N - 1, p, loc=1)  # P(Bin >= k) = 1 - P(Bin <= k-1)
+        # integrate r over chi2_d density
+        val, _ = integrate.quad(lambda r: stats.chi2.pdf(r, d) * (1 - stats.binom.cdf(k - 1, N - 1, stats.ncx2.cdf(d, d, r))),
+                                0, 200 + 20 * d, limit=200)
+        return val
+
+    probs = []
+    for d in d_grid:
+        probs.append(median_prob(d))
+    probs = np.array(probs)
+    # find d* where prob crosses 1/2
+    idx = np.argmin(np.abs(probs - 0.5))
+    return d_grid[idx], probs
+
+
+if __name__ == "__main__":
+    for N in [1000, 10000, 50000]:
+        for k in [1, 5, 10, 20]:
+            dstar, _ = crossover_dimension(N, k)
+            print(f"N={N:6d} k={k:3d} -> d*={dstar}")

databuilders/__init__.py

@@ -0,0 +1 @@
+from .dior import Dior as dior

databuilders/dior.py

@@ -0,0 +1,113 @@
+import os
+import json
+import datasets
+from PIL import Image
+from dataclasses import dataclass
+
+@dataclass
+class DiorConfig(datasets.BuilderConfig):
+    """BuilderConfig for Dior dataset."""
+    pass
+
+
+class Dior(datasets.GeneratorBasedBuilder):
+    """DIOR Dataset."""
+
+    VERSION = datasets.Version("1.0.0")
+
+    BUILDER_CONFIG_CLASS = DiorConfig
+
+    BUILDER_CONFIGS = [
+        DiorConfig(
+            name="default",
+            description="Default configuration for the DIOR dataset.",
+        ),
+    ]
+    
+    DEFAULT_CONFIG_NAME = "default"
+
+    def _info(self):
+
+        return datasets.DatasetInfo(
+            description="The DIOR Dataset.",
+            features=datasets.Features({
+                "image": datasets.Image(),
+                "captions": datasets.Sequence(datasets.Value("string")),
+                "bndboxes": datasets.Array2D(shape=(None, 4), dtype="float32"),
+                "obboxes": datasets.Array2D(shape=(None, 8), dtype="float32"),
+                "dataid": datasets.Value("string")
+            }),
+            homepage="http://www.example.com/",
+            citation="",
+        )
+
+    def _split_generators(self, dl_manager):
+
+        data_files = dl_manager.download_and_extract(self.config.data_files)
+
+        if not data_files or not isinstance(data_files, dict):
+            raise ValueError(
+                "This builder requires you to pass data_files as a dictionary."
+                "for example: data_files={'train': 'path/to/train.jsonl', 'test': 'path/to/test.jsonl'}"
+            )
+        split_generators = []
+
+        for split_name, metadata_list in data_files.items():
+            split_generators.append(
+                datasets.SplitGenerator(
+                    name=split_name,
+                    gen_kwargs={"metadata_list": metadata_list},
+                )
+            )
+        
+        return split_generators
+
+    def _generate_examples(self, metadata_list):
+
+        idx = 0
+
+        for metadata_path in metadata_list:
+            # print(f"--> [Generator] Processing file: {metadata_path}")
+            base_dir = os.path.dirname(os.path.abspath(metadata_path))
+
+            with open(metadata_path, 'r', encoding='utf-8') as f:
+                for i, line in enumerate(f):
+                    try:
+                        data = json.loads(line)
+
+                        absolute_image_path = os.path.join(base_dir, data["file_name"])
+                        # image = Image.open(absolute_image_path).convert("RGB")
+                        
+                        example = {
+                            "image": absolute_image_path,
+                            "captions": data.get("captions", []),
+                            "bndboxes": data.get("bndboxes", []),
+                            "obboxes": data.get("obboxes", []),
+                            "dataid": os.path.splitext(os.path.basename(data["file_name"]))[0]
+                        }
+                        
+                        yield idx, example
+
+                        idx += 1
+
+                    except Exception as e:
+                        print(f"  - Skip Invalid Data ({metadata_path} Line {i}): {e}")
+
+
+if __name__ == "__main__":
+    data_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "data", "DIOR")
+    # data_files = {
+    #     "train": os.path.join(data_dir, "train_meta.jsonl"),
+    #     "test": os.path.join(data_dir, "test_meta.jsonl"),
+    # }
+    data_files = {
+        "train": [os.path.join(data_dir, "train_meta_sample1.jsonl"), 
+                  os.path.join(data_dir, "train_meta_sample2.jsonl")],
+        "test": os.path.join(data_dir, "infer_meta_sample1.jsonl"),
+    }
+
+    builder = Dior(data_files=data_files, config_name="default")
+    builder.download_and_prepare()
+    dataset = builder.as_dataset()
+
+    print(dataset)

datamodules/__init__.py

@@ -0,0 +1,2 @@
+from .loader import Loader
+from .ref_table import RefTable

datamodules/loader.py

@@ -0,0 +1,109 @@
+from . import transforms
+from .ref_table import RefTable
+from utils import get_ckpt_path
+import databuilders
+import datasets
+import torch
+import json
+import os
+
+
+class Loader:
+    def __init__(self, config, image_processor, split='train', logger=None):
+        self.logger = logger
+        self.split, self.phase = split, config.phase
+        self.data_name = config.dataset.name
+        # self.data_files = config.dataset.data_files
+        self.data_files = config.dataset.data_files[self.split][self.phase]
+        self.categories = config.dataset.categories[self.phase]
+        self.image_column, self.caption_column, self.bbox_column, self.obbox_column = config.dataset.column_names
+        self.batch_size = config.training.batch_size
+        self.num_workers = config.training.num_workers
+        self.max_inference_size = config.inference.get('max_inference_size', None)
+        self.novel_sample_dict = None
+
+        self.set_dataset()
+        if self.phase == 'novel':
+            self.k_shot = config.dataset.novel_settings.k_shot
+            self.shuffle_seed = config.dataset.novel_settings.get('shuffle_seed', 42)
+            self.dump_file = os.path.join(get_ckpt_path(config), 'novel_sample_dict.json')
+            if self.split == 'train':
+                self.sample_dataset()
+            elif self.split == 'infer':
+                self.load_novel_sample_dict()
+                self.sample_dataset_infer_phase()
+
+        self.dataset = datasets.concatenate_datasets(self.dataset.values())
+        self.ref_table = RefTable(config, filter_dict=self.novel_sample_dict)
+        self.transform = getattr(transforms, config.dataset.get('transform', 'DefaultTransform'))(config, image_processor, self.split, ref_table=self.ref_table())
+        self.dataset = self.dataset.with_transform(self.transform)
+
+    def set_dataset(self):
+        builder = getattr(databuilders, self.data_name.lower(), None)
+        if builder is None:
+            raise ValueError(f"Unknown dataset: {self.data_name}")
+        builder = builder(data_files=self.data_files)
+        builder.download_and_prepare()
+        self.dataset = builder.as_dataset()
+
+    def sample_dataset(self):
+        ''' The data sample logics for few-shot learning. '''
+        self.novel_sample_dict = {}
+        for category in self.categories:
+            self.dataset[category] = self.dataset[category].shuffle(self.shuffle_seed).select(range(min(self.k_shot, len(self.dataset[category]))))
+            self.novel_sample_dict[category] = list(self.dataset[category]['dataid'])
+
+    def sample_dataset_infer_phase(self):
+        if self.max_inference_size is not None:
+            # self.dataset['default'] = self.dataset['default'].shuffle(self.shuffle_seed).select(range(self.max_inference_size))
+            for category in self.categories:
+                self.dataset[category] = self.dataset[category].shuffle(self.shuffle_seed).select(range(min(self.max_inference_size, len(self.dataset[category]))))
+
+    def collate_fn(self, examples):
+        images = torch.stack([example[self.image_column] for example in examples])
+        images = images.to(memory_format=torch.contiguous_format).float()
+        captions = [example[self.caption_column] for example in examples]
+        bboxes = [example[self.bbox_column] for example in examples]
+        obboxes = [example[self.obbox_column] for example in examples]
+        if isinstance(examples[0]["instances"], list):
+            instances = [example["instances"] for example in examples]
+        else:   
+            instances = torch.stack([example["instances"] for example in examples])
+        dataid = [example["dataid"] for example in examples]
+        # Custom Keys
+        if 'masks' in examples[0].keys():
+            masks = torch.stack([example["masks"] for example in examples])
+            return {self.image_column: images, self.caption_column: captions, self.bbox_column: bboxes, self.obbox_column: obboxes, "instances": instances, "dataid": dataid, "masks": masks}
+        if 'parallels' in examples[0].keys():
+            # parallels = torch.cat([example["parallels"] for example in examples])
+            # images = torch.cat([images, parallels], dim=0)
+            parallels = [example["parallels"] for example in examples]
+            return {self.image_column: images, self.caption_column: captions, self.bbox_column: bboxes, self.obbox_column: obboxes, "instances": instances, "dataid": dataid, "parallels": parallels}
+        return {self.image_column: images, self.caption_column: captions, self.bbox_column: bboxes, self.obbox_column: obboxes, "instances": instances, "dataid": dataid}
+
+    def __call__(self):
+        # for i in range(6):
+        #     _ = self.dataset[i]
+        dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            shuffle=True,
+            collate_fn=self.collate_fn,
+            batch_size=self.batch_size,
+            num_workers=self.num_workers,
+        )
+        return dataloader
+
+    def dump_novel_sample_dict(self):
+        if self.phase == 'novel' and self.logger is not None:
+            self.logger.info(f'Novel Sample Dict: \n{json.dumps(self.novel_sample_dict, indent=2)}', main_process_only=False)
+            self.logger.info(f'Novel Ref Table: \n{json.dumps(self.ref_table("novel"), indent=2)}', main_process_only=False)
+            self.logger.info(f'Dump novel_sample_dict at {self.dump_file}')
+        with open(self.dump_file, 'w') as f:
+            json.dump(self.novel_sample_dict, f, indent=2)
+
+    def load_novel_sample_dict(self):
+        if self.phase == 'novel' and self.logger is not None:
+            self.logger.info(f'Load novel_sample_dict at {self.dump_file}')
+            with open(self.dump_file, 'r') as f:
+                self.novel_sample_dict = json.load(f)
+            self.logger.info(f'Novel Sample Dict: \n{json.dumps(self.novel_sample_dict, indent=2)}', main_process_only=False)

datamodules/ref_table.py

@@ -0,0 +1,76 @@
+import functools
+import imagesize
+import torch
+import os
+import json
+
+
+def singleton(cls):
+    instances = {}
+    def get_instance(*args, **kwargs):
+        if cls not in instances:
+            instances[cls] = cls(*args, **kwargs)
+        return instances[cls]
+    return get_instance
+
+
+@singleton
+class RefTable:
+    def __init__(self, config, filter_dict=None):
+        self.phase = config.phase
+        self.categories_base, self.categories_novel = config.dataset.categories.base, config.dataset.categories.novel
+        self.categories = config.dataset.categories.get('all', None) or (self.categories_novel + self.categories_base)
+        self.image_patch_path = config.dataset.image_patch_path
+        self.filter_dict = filter_dict
+        self.augment = config.dataset.get('ref_augment', False)
+
+        cache_path = os.path.join(self.image_patch_path, "image_sizes.json")
+        if not os.path.exists(cache_path):
+            raise FileNotFoundError(f"File Not Found: {cache_path}")
+        with open(cache_path, 'r') as f:
+            self.img_size_cache = json.load(f)
+        self.ref_table, self.base_ref_table, self.novel_ref_table = self.build_ref_table()
+
+    def build_ref_table(self):
+        ref_table, base_ref_table, novel_ref_table = {}, {}, {}
+        for category in self.categories:
+            category_dir = os.path.join(self.image_patch_path, category)
+            patch_list = os.listdir(category_dir)
+            patch_list = list(filter(lambda s: s.endswith('.jpg'), patch_list))
+
+            # Filter the patch list to avoid data leakage of few-shot learning
+            if self.phase == 'novel' and category in self.categories_novel:
+                assert self.filter_dict is not None
+                patch_list = list(filter(lambda patch_name: patch_name.rsplit('_', 1)[0] in self.filter_dict[category], patch_list))
+                if self.augment:
+                    aug_category_dir = os.path.join(self.image_patch_path, category, 'augmented')
+                    aug_patch_list = os.listdir(aug_category_dir)
+                    aug_patch_list = list(filter(lambda patch_name: patch_name.startswith(tuple(self.filter_dict[category])), aug_patch_list))
+                    aug_patch_list = list(map(lambda patch_name: f"augmented/{patch_name}", aug_patch_list))
+                    patch_list += aug_patch_list
+
+            def get_size(img_name):
+                return self.img_size_cache.get(img_name, [0, 1])
+
+            patch_list = sorted(patch_list, key = lambda img: get_size(img)[0] * get_size(img)[1], reverse=True)
+
+            # Only build novel ref table at novel phase
+            if self.phase == 'novel' and category in self.categories_novel:
+                novel_ref_table[category] = {
+                    img: get_size(img)[0] / get_size(img)[1] for img in patch_list[:200]
+                }
+            elif category in self.categories_base:
+                base_ref_table[category] = {
+                    img: get_size(img)[0] / get_size(img)[1] for img in patch_list[:200]
+                }
+
+        ref_table = novel_ref_table | base_ref_table
+        return ref_table, base_ref_table, novel_ref_table
+
+    def __call__(self, phase=None):
+        if phase is None:
+            return self.ref_table
+        if phase == 'base':
+            return self.base_ref_table
+        if phase == 'novel':
+            return self.novel_ref_table

datamodules/transforms/__init__.py

@@ -0,0 +1 @@
+from .layout_transform import LayoutTransform

datamodules/transforms/layout_transform.py

@@ -0,0 +1,132 @@
+import albumentations as A
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import functools
+import imagesize
+import torch
+import cv2
+import os
+
+
+class LayoutTransform:
+    def __init__(self, config, image_processor, split, ref_table=None, filter_dict=None):
+        self.split, self.phase = split, config.phase
+        if split == 'train' and self.phase == 'novel':
+            self.image_transforms = A.Compose([
+                A.OneOf([
+                    A.RandomSizedBBoxSafeCrop(height=config.dataset.resolution, width=config.dataset.resolution, erosion_rate=0.0, interpolation=cv2.INTER_CUBIC, p=0.3),
+                    A.Resize(height=config.dataset.resolution, width=config.dataset.resolution, interpolation=cv2.INTER_CUBIC, p=0.7),
+                ], p=1.0),
+                A.Normalize(mean=[0.5], std=[0.5]),
+                A.pytorch.ToTensorV2(),
+            ], bbox_params=A.BboxParams(format='albumentations', label_fields=['labels'], min_area=0, min_visibility=0.0))
+        elif split == 'infer' or self.phase == 'base':
+            self.image_transforms = A.Compose([
+                A.Resize(config.dataset.resolution, config.dataset.resolution),
+                A.Normalize(mean=[0], std=[1]),
+                A.pytorch.ToTensorV2(),
+            ])
+        else:
+            raise ValueError("Invalid mode for Transform.")
+        self.image_patch_path = config.dataset.image_patch_path
+        self.ref_resolution = config.dataset.ref_resolution
+        self.image_column, self.caption_column, self.bbox_column, self.obbox_column = config.dataset.column_names
+        self.image_processor = image_processor
+        if ref_table is not None:
+            self.ref_table = ref_table
+        else:
+            self.categories = config.dataset.categories[self.phase]
+            self.filter_dict = filter_dict
+            self.ref_table = self.build_ref_table()
+        self.k_shot = config.dataset.novel_settings.k_shot
+        self.top_k = 1
+
+    def build_ref_table(self):
+        ref_table = {}
+        for category in self.categories:
+            category_dir = os.path.join(self.image_patch_path, category)
+            patch_list = os.listdir(category_dir)
+            # Filter the patch list to avoid data leakage of few-shot learning
+            if self.phase == 'novel':
+                assert self.filter_dict is not None
+                patch_list = list(filter(lambda patch_name: patch_name.split('_')[0] in self.filter_dict[category], patch_list))
+            patch_list = sorted(patch_list, key = lambda img: functools.reduce(lambda x, y: x*y, imagesize.get(os.path.join(category_dir, img))), reverse=True)
+            ref_table[category] = {img: functools.reduce(lambda x, y: x/y, imagesize.get(os.path.join(category_dir, img))) for img in patch_list[:200]}
+        return ref_table
+
+    @staticmethod
+    def find_nearest(array, value):
+        array = np.asarray(array)
+        idx = (np.abs(array/value - 1)).argmin()
+        return idx
+    
+    @staticmethod
+    def find_k_nearest(array, value, k=1):
+        array = np.asarray(array)
+        dist = np.abs(array / value - 1)
+        idxs = np.argsort(dist)[:k]
+        return idxs
+    
+    def get_instances(self, examples):
+        instances = []
+        for index, (caption, bboxes) in enumerate(zip(examples[self.caption_column], examples[self.bbox_column])):
+            categories = caption[1:]
+            instances_per_example = []
+            for name, bbox in zip(categories, bboxes):
+                if name == '':
+                    instances_per_example.append(torch.zeros([self.top_k, 3, self.ref_resolution, self.ref_resolution]))
+                else:
+                    value = (bbox[2] - bbox[0]) / max(bbox[3] - bbox[1], 1e-8)
+                    chosen_idxs = self.find_k_nearest(list(self.ref_table[name].values()), value, k=self.top_k)
+                    instances_per_bbox = []
+                    for idx in chosen_idxs:
+                        chosen_file = list(self.ref_table[name].keys())[idx]
+                        img = Image.open(os.path.join(self.image_patch_path, name, chosen_file)).convert('RGB')
+                        img = self.image_processor(images=img, return_tensors="pt")['pixel_values'].squeeze(0)
+                        instances_per_bbox.append(img)
+                    instances_per_example.append(torch.stack(instances_per_bbox))
+            instances.append(torch.stack(instances_per_example))
+        return instances
+    
+    def train_transform(self, examples):
+        images, bboxes, obboxes = examples[self.image_column], examples[self.bbox_column], examples[self.obbox_column]
+        global_prompt = [caption[:1] for caption in examples[self.caption_column]]
+        captions = [caption[1:] for caption in examples[self.caption_column]]
+
+        new_images, new_bboxes, new_obboxes, new_captions = [], [], [], []
+        for i in range(len(images)):
+            num_instances = sum(bool(s) for s in captions[i])
+            bboxes_i = bboxes[i][:num_instances]
+            captions_i = captions[i][:num_instances]
+            transformed = self.image_transforms(image=images[i], bboxes=bboxes_i, labels=captions_i)
+            transformed["obboxes"] = []
+            for xmin, ymin, xmax, ymax in transformed["bboxes"]:
+                transformed["obboxes"].append([xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax])
+            for _ in range(len(captions[i]) - num_instances):
+                transformed["bboxes"].append([0,0,0,0])
+                transformed["obboxes"].append([0,0,0,0,0,0,0,0])
+                transformed["labels"].append("")
+            new_images.append(transformed["image"])
+            new_bboxes.append(transformed["bboxes"])
+            new_obboxes.append(transformed["obboxes"])
+            new_captions.append(global_prompt[i] + transformed["labels"])
+        examples[self.image_column] = new_images
+        examples[self.bbox_column] = new_bboxes
+        examples[self.obbox_column] = new_obboxes
+        examples[self.caption_column] = new_captions
+        return examples
+
+    def infer_transform(self, examples):
+        examples[self.image_column] = [self.image_transforms(image=image)['image'] for image in examples[self.image_column]]
+        return examples
+
+    def __call__(self, examples):
+        # print("Received keys:", examples.keys())
+        examples[self.image_column] = [np.array(image.convert("RGB")) for image in examples[self.image_column]]
+        if self.split == "train":
+            examples = self.train_transform(examples)
+        elif self.split == "infer":
+            examples = self.infer_transform(examples)
+        examples["instances"] = self.get_instances(examples)
+        return examples

fonts/Rainbow-Party-2.ttf

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

gicdm_core.py

@@ -0,0 +1,216 @@
+"""Core GICDM implementation reproducing the paper arXiv:2602.16449.
+
+Implements:
+  - ICDM iterative hubness reduction (Section 3 / Algorithm in paper lines 485-498)
+  - GICDM out-of-sample generated-point scaling (Eq. 1, Algorithm 1)
+  - Hubness statistics h5_1(1%) and A5 (Table 4)
+  - Clipped Density / Clipped Coverage fidelity & coverage metrics
+  - Crossover dimension d* (Proposition 5.3)
+  - Raisa et al. (2025) synthetic benchmark scenarios
+"""
+import numpy as np
+from scipy import stats
+from scipy.special import ncfdtr, ncfdtri
+
+
+# ----------------------------------------------------------------------------
+# Distances & k-nearest-neighbour helpers
+# ----------------------------------------------------------------------------
+def pairwise_sq_dists(X, Y=None):
+    if Y is None:
+        Y = X
+    Xn = np.sum(X ** 2, axis=1, keepdims=True)
+    Yn = np.sum(Y ** 2, axis=1, keepdims=True)
+    D2 = Xn + Yn.T - 2.0 * (X @ Y.T)
+    np.maximum(D2, 0, out=D2)
+    return np.sqrt(D2)
+
+
+def knn_distances(D, k, self_included=True):
+    """Return (N, k) sorted distances to k nearest neighbours.
+
+    If self_included, the 0-th distance (0) is the point itself.
+    """
+    N = D.shape[0]
+    if self_included:
+        idx = np.argpartition(D, k, axis=1)[:, :k]
+    else:
+        # exclude self (diagonal)
+        Dc = D + np.eye(N) * 1e18
+        idx = np.argpartition(Dc, k, axis=1)[:, :k]
+    idx = idx[np.arange(N)[:, None], np.argsort(D[np.arange(N)[:, None], idx], axis=1)]
+    return D[np.arange(N)[:, None], idx]
+
+
+# ----------------------------------------------------------------------------
+# ICDM (Iterative Contextual Dissimilarity Measure)
+# Matches official implementation: metrics/hubness_processor/hubness_reduction_methods.py
+# ----------------------------------------------------------------------------
+def _average_k_dist(D, k):
+    """Average of the k nearest distances (including self at distance 0),
+    i.e. sum of (k+1) smallest distances / k."""
+    k_dists = np.partition(D, k, axis=1)[:, : k + 1]
+    return k_dists.sum(axis=1) / k
+
+
+def icdm_scaling(D, K, n_iter=10, return_mu=False):
+    """ICDM scaling factors delta_i, matching official `icdm_delta_low_memory`.
+
+    Iteratively applies NICDM:
+        d_{ij}  <- d_{ij} / (sqrt(r_i) sqrt(r_j))   with r_i = average of k nearest dists (k=K)
+        delta_i <- delta_i / sqrt(r_i)
+    Final: secondary dissimilarity d^T_{ij} = d_{ij} delta_i delta_j.
+    Returns delta_i (length N) and optionally final average-neighbour distances.
+    """
+    N = D.shape[0]
+    d = D.copy()
+    deltas = np.ones(N)
+    for _ in range(n_iter):
+        r = _average_k_dist(d, K)
+        sqrt_r = np.maximum(np.sqrt(r), 1e-12)
+        d = d / (sqrt_r[:, None] * sqrt_r[None, :])
+        deltas = deltas / sqrt_r
+    if return_mu:
+        mu_final = _average_k_dist(d, K)
+        return deltas, mu_final
+    return deltas
+
+
+# ----------------------------------------------------------------------------
+# GICDM (Algorithm 1)
+# ----------------------------------------------------------------------------
+def gicdm(Xr, Xg, K1, K2, q=0.95, n_iter=10, return_info=False):
+    """Generative ICDM (Algorithm 1), matching the official implementation.
+
+    Xr : (N, d) real points
+    Xg : (M, d) generated points
+    K1, K2 : two filter scales (K2 = 10*K1 in the paper). For each scale k,
+             ICDM is applied with neighbourhood size 2k (paper: GICDM K = 2k).
+    Returns:
+        D_final : (M, N) GICDM real-to-generated dissimilarity matrix
+        keep    : (M,) boolean mask of generated points passing multi-scale filter
+        Drr_gicdm : (N, N) GICDM real-to-real dissimilarity matrix
+    """
+    Drr = pairwise_sq_dists(Xr)
+    Drg = pairwise_sq_dists(Xg, Xr)  # (M, N): generated-to-real
+    N, M = Xr.shape[0], Xg.shape[0]
+
+    keep = np.ones(M, dtype=bool)
+    info = {}
+
+    for k in (K1, K2):
+        # ICDM with neighbourhood 2k
+        delta_r = icdm_scaling(Drr, K=2 * k, n_iter=n_iter)
+
+        # --- real filter (Algorithm 1 lines 3-6) ---
+        # k nearest real neighbours (exclude self)
+        nn_r = np.argpartition(Drr, k, axis=1)[:, : k + 1]
+        is_self = nn_r == np.arange(N)[:, None]
+        no_self = ~is_self.any(axis=1)
+        is_self[no_self, -1] = True
+        nn_r = nn_r[~is_self].reshape(N, k)
+        real_avg_delta = delta_r[nn_r].mean(axis=1)
+        r_ri = np.abs(real_avg_delta - delta_r) / real_avg_delta
+        T_k = np.quantile(r_ri, q)
+
+        # --- generated filter (Algorithm 1 lines 9-14) ---
+        # k+1 nearest real neighbours of each generated point
+        nn_g = np.argpartition(Drg.T, k, axis=1)[:, : k + 1]
+        # r_synthetic = average of (delta_real_nn * d_orig) over k+1 neighbours
+        d_g_nn = Drg[np.arange(M)[:, None], nn_g]
+        delta_r_g_nn = delta_r[nn_g]
+        r_synthetic = (delta_r_g_nn * d_g_nn).sum(axis=1) / (k + 1)
+        delta_g = 1.0 / r_synthetic  # Eq. (1) with mu_bar absorbed
+        synth_avg_delta = delta_r_g_nn.mean(axis=1)
+        r_gj = np.abs(synth_avg_delta - delta_g) / synth_avg_delta
+        keep &= (r_gj <= T_k)
+
+        info[k] = dict(T_k=float(T_k), delta_g=delta_g)
+
+    # final GICDM dissimilarities (Algorithm 1 line 17)
+    delta_g_K1 = info[K1]['delta_g']
+    delta_r_K1 = icdm_scaling(Drr, K=2 * K1, n_iter=n_iter)
+    D_final = Drg * (delta_r_K1[None, :] * delta_g_K1[None, :])
+    Drr_gicdm = Drr * np.outer(delta_r_K1, delta_r_K1)
+
+    if return_info:
+        return D_final, keep, info, (delta_r_K1, delta_g_K1)
+    return D_final, keep, Drr_gicdm
+
+
+# ----------------------------------------------------------------------------
+# Hubness statistics
+# ----------------------------------------------------------------------------
+def k_occurrence(D, k=5):
+    """O_k(x_i) = number of points for which x_i is among their k NN (excl self)."""
+    N = D.shape[0]
+    Dc = D + np.eye(N) * 1e18
+    knn = np.argsort(Dc, axis=1)[:, :k]
+    occ = np.zeros(N, dtype=int)
+    for j in range(k):
+        occ[knn[:, j]] += 1
+    return occ
+
+
+def hubness_stats(D, k=5, q=0.01):
+    """Return h5_1(1%) and A5 (proportion of antihubs)."""
+    occ = k_occurrence(D, k)
+    mean_occ = occ.mean()
+    n = len(occ)
+    topq = max(1, int(np.floor(q * n)))
+    top_vals = np.sort(occ)[::-1][:topq]
+    h5 = top_vals.mean() / mean_occ if mean_occ > 0 else np.nan
+    A5 = float(np.mean(occ == 0))
+    return float(h5), A5
+
+
+# ----------------------------------------------------------------------------
+# Clipped Density / Clipped Coverage (Salvy et al. 2026)
+# ----------------------------------------------------------------------------
+def clipped_density(Xr, Xg, k=5, dissim=None, keep=None):
+    """Clipped Density fidelity metric.
+
+    For each generated point, distance to its k-th real NN; threshold = distance
+    from each real point to its k-th real NN (clip). Score averages clip term.
+    If dissim ('gicdm') is provided, use GICDM dissimilarities instead of raw dist.
+    keep: boolean mask of generated points to include (filtered-out points get 0).
+    """
+    Drg = pairwise_sq_dists(Xg, Xr)  # (M,N)
+    Drr = pairwise_sq_dists(Xr)
+    if dissim is None:
+        d_rg = Drg
+        d_rr = Drr
+    else:
+        d_rg, d_rr = dissim
+    # k-th NN distance for each real point (in its own set)
+    kth_real = np.sort(d_rr + np.eye(len(Xr)) * 1e18, axis=1)[:, k - 1]
+    kth_gen = np.sort(d_rg, axis=1)[:, k - 1]
+    # clip each generated point's distance at its matched real threshold
+    thresh = kth_real[np.argmin(d_rg, axis=1)]
+    clip = np.clip(kth_gen / thresh, 0, 1)
+    if keep is not None:
+        clip = clip * keep.astype(float)  # filtered points -> fidelity 0
+    return float(np.mean(clip))
+
+
+def clipped_coverage(Xr, Xg, k=5, dissim=None, keep=None):
+    """Clipped Coverage: for each real point, does a kept generated point fall
+    within its k-th NN threshold?"""
+    Drg = pairwise_sq_dists(Xg, Xr)
+    Drr = pairwise_sq_dists(Xr)
+    if dissim is None:
+        d_rg = Drg
+        d_rr = Drr
+    else:
+        d_rg, d_rr = dissim
+    kth_real = np.sort(d_rr + np.eye(len(Xr)) * 1e18, axis=1)[:, k - 1]
+    min_d = d_rg.min(axis=0)
+    covered = (min_d <= kth_real).astype(float)
+    if keep is not None:
+        # a generated point only contributes if it is kept
+        d_rg_k = d_rg[keep, :]
+        if d_rg_k.shape[0] == 0:
+            return 0.0
+        min_d2 = d_rg_k.min(axis=0)
+        covered = (min_d2 <= kth_real).astype(float)
+    return float(np.mean(covered))

infer.sh

@@ -0,0 +1,174 @@
+#!/usr/bin/env bash
+set -e
+
+CONFIGS=()
+SEEDS=()
+CKPTS=()
+RUN_IDS=()
+GPU_IDS=""
+METASEED=""
+NUM_SEED=-1
+META_START=0
+MAX_INFER_SIZE=""
+K_SHOTS=()
+
+while [[ $# -gt 0 ]]; do
+    case $1 in
+        --config)
+            shift
+            CONFIGS=($1)
+            ;;
+        --seed)
+            shift
+            SEEDS=($1)
+            ;;
+        --metaseed)
+            shift
+            METASEED=$1
+            ;;
+        --num_seed)
+            shift
+            NUM_SEED=$1
+            ;;
+        --meta_start)
+            shift
+            META_START=$1
+            ;;
+        --ckpt)
+            shift
+            CKPTS=($1)
+            ;;
+        --run_id)
+            shift
+            RUN_IDS=($1)
+            ;;
+        --k_shot)
+            shift
+            K_SHOTS=($1)
+            ;;
+        --gpu_ids)
+            shift
+            GPU_IDS=$1
+            ;;
+        --max_infer_size)
+            shift
+            MAX_INFER_SIZE=$1
+            ;;
+        *)
+            echo "Unknown argument: $1"
+            exit 1
+            ;;
+    esac
+    shift
+done
+
+if [[ ${#CONFIGS[@]} -eq 0 ]] \
+ || [[ ${#CKPTS[@]} -eq 0 ]] \
+ || [[ -z "$GPU_IDS" ]] \
+ || [[ ${#K_SHOTS[@]} -eq 0 ]] \
+ || [[ ${#RUN_IDS[@]} -eq 0 ]]; then
+    echo "Error: missing required arguments."
+    echo "You must provide: --config, --gpu_ids, --run_id, --k_shot"
+    echo "And one of: --seed OR --metaseed"
+    exit 1
+fi
+
+if [[ -z "$METASEED" && ${#SEEDS[@]} -eq 0 ]]; then
+    echo "Error: either --seed or --metaseed must be provided."
+    exit 1
+fi
+
+IFS=',' read -r -a GPU_ARRAY <<< "$GPU_IDS"
+NUM_PROCESSES=${#GPU_ARRAY[@]}
+
+if [[ -n "$METASEED" ]]; then
+    if [[ $NUM_SEED -lt 0 ]]; then
+        echo "Error: you must set --num_seed when using --metaseed."
+        exit 1
+    fi
+
+    echo "[INFO] Using metaseed=$METASEED"
+    echo "[INFO] Will use seeds index range [$META_START, $NUM_SEED)"
+
+    mapfile -t ALL_SEEDS < <(
+        shuf -i 0-9999 \
+            --random-source=<(awk -v s="$METASEED" 'BEGIN { while (1) printf "%s", s }') \
+            | head -n $NUM_SEED
+    )
+
+    # SEEDS=("${ALL_SEEDS[@]:META_START:NUM_SEED-META_START}")
+    SEEDS=("${ALL_SEEDS[@]:$META_START:$((NUM_SEED - META_START))}")
+
+    echo "[INFO] Total Seeds Generated: ${#ALL_SEEDS[@]}"
+    echo "[INFO] Using Seeds: ${SEEDS[@]}"
+fi
+
+echo "CONFIGS:   ${CONFIGS[@]}"
+echo "SEEDS:     ${SEEDS[@]}"
+echo "RUN_IDS:    ${RUN_IDS[@]}"
+echo "K_SHOTS:    ${K_SHOTS[@]}"
+echo "CKPTS:     ${CKPTS[@]}"
+echo "GPU_IDS:   ${GPU_IDS}"
+echo "NUM_PROCESSES: $NUM_PROCESSES"
+echo "META_START: $META_START"
+echo "NUM_SEED: $NUM_SEED"
+echo "Actual inference episodes: ${#SEEDS[@]}"
+echo "MAX_INFER_SIZE: ${MAX_INFER_SIZE:-"(unset)"}"
+
+EXTRA_ARG=""
+if [[ -n "$MAX_INFER_SIZE" ]]; then
+    EXTRA_ARG="-M $MAX_INFER_SIZE"
+fi
+
+run_with_retry(){
+    local cmd=("$@")
+    local max=3
+    local attempt=1
+
+    while (( attempt <= max )); do
+        echo "Attempt $attempt/$max: ${cmd[*]}"
+        
+        set +e
+        "${cmd[@]}"
+        status=$?
+        set -e
+
+        if [[ $status -eq 0 ]]; then
+            return 0
+        fi
+
+        ((attempt++))
+        sleep 2
+    done
+
+    return 1
+}
+
+for k_shot in "${K_SHOTS[@]}"; do
+    for run_id in "${RUN_IDS[@]}"; do
+        for config in "${CONFIGS[@]}"; do
+            for seed in "${SEEDS[@]}"; do
+                for ckpt in "${CKPTS[@]}"; do
+                    echo "config: $config, seed: $seed, ckpt: $ckpt, run: $run_id"
+
+                    max_wait=12
+                    ckpt_path="./ckpt/$config/novel/run-$run_id/$k_shot-shot/shuffle_seed-$seed/checkpoint-$ckpt"
+                    while [[ ! -s "$ckpt_path" ]]; do
+                        echo "waiting for ckpt: $ckpt_path"
+                        sleep 300
+                        ((max_wait--))
+                    done
+
+                    if [[ ! -s "$ckpt_path" ]]; then
+                        echo "[ERROR] ckpt not found after waiting, skip."
+                        continue
+                    fi
+
+                    run_with_retry \
+                    accelerate launch --multi_gpu --gpu_ids $GPU_IDS --num_processes $NUM_PROCESSES main.py \
+                    --config ./configs/$config.yaml -m infer -p novel -s $seed -k $k_shot -r $run_id -c $ckpt $EXTRA_ARG
+                done
+            done
+        done
+    done
+done

main.py

@@ -0,0 +1,12 @@
+from utils import load_config
+import variants
+
+if __name__ == '__main__':
+    config = load_config()
+
+    if config.mode == 'train':
+        entry = getattr(variants.train, config.training.get('entry', 'train_dsp'))
+    elif config.mode == 'infer':
+        entry = getattr(variants.infer, config.inference.get('entry', 'infer_dsp'))
+
+    entry(config)

models/__init__.py

@@ -0,0 +1 @@
+from . import dsp

models/dsp/CAM/__init__.py

@@ -0,0 +1 @@
+from .cam_generator import CAMGenerator

models/dsp/CAM/cam_generator.py

@@ -0,0 +1,133 @@
+from .pytorch_grad_cam import GradCAM
+from .pytorch_grad_cam.utils.image import scale_cam_image
+from . import clip
+from .utils import reshape_transform, zeroshot_classifier, ClipOutputTarget, scoremap2bbox
+
+from torchvision import transforms
+import torch
+import numpy as np
+import cv2
+
+class CAMGenerator:
+    def __init__(self, categories, clip_path):
+        # self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        self.clip_path = clip_path
+        self.clip_model, _ = clip.load(self.clip_path, device="cpu")
+        self.target_layers = [self.clip_model.visual.transformer.resblocks[-1].ln_1]
+        self.cam = GradCAM(model=self.clip_model, target_layers=self.target_layers, reshape_transform=reshape_transform, use_cuda=True)
+
+        self.categories = categories
+        # if categories is None:
+        #     self.categories = ['vehicle', 'baseballfield', 'groundtrackfield', 'windmill', 'bridge', 'overpass', 'ship', 'airplane', 'tenniscourt', 'airport',
+        #                     'expressway-service-area', 'basketballcourt', 'stadium', 'storagetank', 'chimney', 'dam', 'expressway-toll-station', 'golffield', 'trainstation', 'harbor']
+        self.background_categories = ['ground','land','grass','tree','building','wall','sky','lake','water','river','sea','railway','railroad','keyboard','helmet', 
+                                      'cloud','house','mountain','ocean','road','rock','street','valley','bridge','sign',]
+
+        self.normalize = transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711))
+
+    def _prepare(self):
+        self.bg_text_features = zeroshot_classifier(self.background_categories, ['a clean origami {}.'], self.clip_model, self.device)
+        self.fg_text_features = zeroshot_classifier(self.categories, ['a clean origami {}.'], self.clip_model, self.device)
+
+    def to(self, device, dtype):
+        self.device = device
+        self.clip_model.to(device)
+        self.cam.set_device(device)
+        self._prepare()
+
+    def re_normalize(self, image):
+        # image = (image / 2) + 0.5
+        image = self.normalize(image)
+        return image
+    
+    def get_label_list(self, captions):
+        label_list = []
+        label_id_list = []
+        for caption in captions:
+            if caption in self.categories and caption not in label_list:
+                label_list.append(caption)
+                label_id_list.append(self.categories.index(caption))
+        return label_list, label_id_list
+    
+    def get_label_list_with_bboxes(self, captions, bboxes):
+        label_list, label_id_list, bboxes_list = [], [], []
+        for caption, bbox in zip(captions, bboxes):
+            if caption not in self.categories:
+                continue
+            if caption not in label_list:
+                label_list.append(caption)
+                label_id_list.append(self.categories.index(caption))
+                bboxes_list.append([bbox])
+            else:
+                bboxes_list[label_list.index(caption)].append(bbox)
+        return label_list, label_id_list, bboxes_list
+
+    def __call__(self, image, captions, bboxes, gt_bboxes_only=False):
+        image = self.re_normalize(image)
+        # label_list, label_id_list = self.get_label_list(captions[0][1:])
+        label_list, label_id_list, bboxes_list = self.get_label_list_with_bboxes(captions[0][1:], bboxes[0])
+        h, w = image.shape[-2], image.shape[-1]
+        image_features, attn_weight_list = self.clip_model.encode_image(image, h, w)
+
+        bg_features_temp = self.bg_text_features
+        fg_features_temp = self.fg_text_features[label_id_list]
+        text_features_temp = torch.cat([fg_features_temp, bg_features_temp], dim=0)
+        input_tensor = [image_features, text_features_temp, h, w]
+
+        keys, refined_cam_list = [], []
+        for idx, (label, bbox) in enumerate(zip(label_list, bboxes_list)):
+            keys.append(self.categories.index(label))
+            targets = [ClipOutputTarget(label_list.index(label))]
+            grayscale_cam, logits_per_image, attn_weight_last = self.cam(input_tensor=input_tensor, targets=targets, target_size=None)
+            grayscale_cam = grayscale_cam[0, :] # [32, 32]
+            # grayscale_cam_highres = cv2.resize(grayscale_cam, (ori_width, ori_height))
+
+            if idx == 0:
+                attn_weight_list.append(attn_weight_last)
+                attn_weight = [aw[:, 1:, 1:] for aw in attn_weight_list]  # (b, hxw, hxw)
+                attn_weight = torch.stack(attn_weight, dim=0)[-8:] # [8, 1, 1024, 1024]
+                attn_weight = torch.mean(attn_weight, dim=0) # [1, 1024, 1024]
+                # attn_weight = attn_weight[0].detach() # [1024, 1024] # original detach
+                attn_weight = attn_weight[0] #.detach() # [1024, 1024]
+            attn_weight = attn_weight.float()
+
+            gt_box, gt_cnt = (np.array(bbox) * [grayscale_cam.shape[1], grayscale_cam.shape[0], grayscale_cam.shape[1], grayscale_cam.shape[0]]).astype(int), len(bbox)
+            if gt_bboxes_only:
+                box, cnt = gt_box, gt_cnt
+            else:
+                box, cnt = scoremap2bbox(scoremap=grayscale_cam.cpu().data.numpy(), threshold=0.4, multi_contour_eval=True)
+                box, cnt = np.concatenate([box, gt_box], axis=0), cnt + gt_cnt
+            aff_mask = torch.zeros_like(grayscale_cam)
+            for i_ in range(cnt):
+                x0_, y0_, x1_, y1_ = box[i_]
+                aff_mask[y0_:y1_, x0_:x1_] = 1
+            aff_mask = aff_mask.view(1, grayscale_cam.shape[0] * grayscale_cam.shape[1])
+
+            aff_mat = attn_weight
+            trans_mat = aff_mat / torch.sum(aff_mat, dim=0, keepdim=True)
+            trans_mat = trans_mat / torch.sum(trans_mat, dim=1, keepdim=True)
+            for _ in range(2):
+                trans_mat = trans_mat / torch.sum(trans_mat, dim=0, keepdim=True)
+                trans_mat = trans_mat / torch.sum(trans_mat, dim=1, keepdim=True)
+            trans_mat = (trans_mat + trans_mat.transpose(1, 0)) / 2
+            for _ in range(1):
+                trans_mat = torch.matmul(trans_mat, trans_mat)
+
+            trans_mat = trans_mat * aff_mask
+
+            cam_to_refine = grayscale_cam.view(-1, 1)
+            cam_refined = torch.matmul(trans_mat, cam_to_refine).reshape(h //16, w // 16)
+            cam_refined = cam_refined - cam_refined.min()
+            cam_refined = cam_refined / (cam_refined.max() + 1e-7)
+            refined_cam_list.append(cam_refined)
+
+        keys = torch.tensor(keys)
+        refined_cams = torch.stack(refined_cam_list, dim=0)
+
+        return refined_cams, keys
+
+
+if __name__ == '__main__':
+    cam = CAMGenerator()
+    

models/dsp/CAM/clip/__init__.py

@@ -0,0 +1 @@
+from .clip import *

models/dsp/CAM/clip/bpe_simple_vocab_16e6.txt.gz

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

models/dsp/CAM/clip/clip.py

@@ -0,0 +1,245 @@
+import hashlib
+import os
+import urllib
+import warnings
+from typing import Any, Union, List
+from packaging import version
+
+import torch
+from PIL import Image
+from torchvision.transforms import Compose, Resize, CenterCrop, ToTensor, Normalize
+from tqdm import tqdm
+
+from .model import build_model
+from .simple_tokenizer import SimpleTokenizer as _Tokenizer
+from collections import OrderedDict
+
+try:
+    from torchvision.transforms import InterpolationMode
+    BICUBIC = InterpolationMode.BICUBIC
+except ImportError:
+    BICUBIC = Image.BICUBIC
+
+
+if version.parse(torch.__version__) < version.parse("1.7.1"):
+    warnings.warn("PyTorch version 1.7.1 or higher is recommended")
+
+
+__all__ = ["available_models", "load", "tokenize"]
+_tokenizer = _Tokenizer()
+
+_MODELS = {
+    "RN50": "https://openaipublic.azureedge.net/clip/models/afeb0e10f9e5a86da6080e35cf09123aca3b358a0c3e3b6c78a7b63bc04b6762/RN50.pt",
+    "RN101": "https://openaipublic.azureedge.net/clip/models/8fa8567bab74a42d41c5915025a8e4538c3bdbe8804a470a72f30b0d94fab599/RN101.pt",
+    "RN50x4": "https://openaipublic.azureedge.net/clip/models/7e526bd135e493cef0776de27d5f42653e6b4c8bf9e0f653bb11773263205fdd/RN50x4.pt",
+    "RN50x16": "https://openaipublic.azureedge.net/clip/models/52378b407f34354e150460fe41077663dd5b39c54cd0bfd2b27167a4a06ec9aa/RN50x16.pt",
+    "RN50x64": "https://openaipublic.azureedge.net/clip/models/be1cfb55d75a9666199fb2206c106743da0f6468c9d327f3e0d0a543a9919d9c/RN50x64.pt",
+    "ViT-B/32": "https://openaipublic.azureedge.net/clip/models/40d365715913c9da98579312b702a82c18be219cc2a73407c4526f58eba950af/ViT-B-32.pt",
+    "ViT-B/16": "https://openaipublic.azureedge.net/clip/models/5806e77cd80f8b59890b7e101eabd078d9fb84e6937f9e85e4ecb61988df416f/ViT-B-16.pt",
+    "ViT-L/14": "https://openaipublic.azureedge.net/clip/models/b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836/ViT-L-14.pt",
+    "ViT-L/14@336px": "https://openaipublic.azureedge.net/clip/models/3035c92b350959924f9f00213499208652fc7ea050643e8b385c2dac08641f02/ViT-L-14-336px.pt",
+}
+
+
+def _download(url: str, root: str):
+    os.makedirs(root, exist_ok=True)
+    filename = os.path.basename(url)
+
+    expected_sha256 = url.split("/")[-2]
+    download_target = os.path.join(root, filename)
+
+    if os.path.exists(download_target) and not os.path.isfile(download_target):
+        raise RuntimeError(f"{download_target} exists and is not a regular file")
+
+    if os.path.isfile(download_target):
+        if hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256:
+            return download_target
+        else:
+            warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
+
+    with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
+        with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
+            while True:
+                buffer = source.read(8192)
+                if not buffer:
+                    break
+
+                output.write(buffer)
+                loop.update(len(buffer))
+
+    if hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256:
+        raise RuntimeError(f"Model has been downloaded but the SHA256 checksum does not not match")
+
+    return download_target
+
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def _transform(n_px):
+    return Compose([
+        Resize(n_px, interpolation=BICUBIC),
+        CenterCrop(n_px),
+        _convert_image_to_rgb,
+        ToTensor(),
+        Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+    ])
+
+
+def available_models() -> List[str]:
+    """Returns the names of available CLIP models"""
+    return list(_MODELS.keys())
+
+
+def load(name: str, device: Union[str, torch.device] = "cuda" if torch.cuda.is_available() else "cpu", jit: bool = False, download_root: str = None):
+    """Load a CLIP model
+
+    Parameters
+    ----------
+    name : str
+        A model name listed by `clip.available_models()`, or the path to a model checkpoint containing the state_dict
+
+    device : Union[str, torch.device]
+        The device to put the loaded model
+
+    jit : bool
+        Whether to load the optimized JIT model or more hackable non-JIT model (default).
+
+    download_root: str
+        path to download the model files; by default, it uses "~/.cache/clip"
+
+    Returns
+    -------
+    model : torch.nn.Module
+        The CLIP model
+
+    preprocess : Callable[[PIL.Image], torch.Tensor]
+        A torchvision transform that converts a PIL image into a tensor that the returned model can take as its input
+    """
+    if name in _MODELS:
+        model_path = _download(_MODELS[name], download_root or os.path.expanduser("~/.cache/clip"))
+    elif os.path.isfile(name):
+        model_path = name
+    else:
+        raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
+
+    with open(model_path, 'rb') as opened_file:
+        try:
+            # loading JIT archive
+            model = torch.jit.load(opened_file, map_location=device if jit else "cpu").eval()
+            state_dict = None
+        except RuntimeError:
+            # loading saved state dict
+            if jit:
+                warnings.warn(f"File {model_path} is not a JIT archive. Loading as a state dict instead")
+                jit = False
+            if 'RN50' in model_path:
+                state_dict = torch.load(opened_file, map_location="cpu")
+            else:
+                state_dict0 = torch.load(model_path, map_location="cpu")
+                state_dict = OrderedDict()
+                for k in state_dict0.keys():
+                    state_dict[k.replace('module.', '')] = state_dict0[k]
+
+
+    if not jit:
+        model = build_model(state_dict or model.state_dict()).to(device)
+        if str(device) == "cpu":
+            model.float()
+        return model, _transform(model.visual.input_resolution)
+
+    # patch the device names
+    device_holder = torch.jit.trace(lambda: torch.ones([]).to(torch.device(device)), example_inputs=[])
+    device_node = [n for n in device_holder.graph.findAllNodes("prim::Constant") if "Device" in repr(n)][-1]
+
+    def patch_device(module):
+        try:
+            graphs = [module.graph] if hasattr(module, "graph") else []
+        except RuntimeError:
+            graphs = []
+
+        if hasattr(module, "forward1"):
+            graphs.append(module.forward1.graph)
+
+        for graph in graphs:
+            for node in graph.findAllNodes("prim::Constant"):
+                if "value" in node.attributeNames() and str(node["value"]).startswith("cuda"):
+                    node.copyAttributes(device_node)
+
+    model.apply(patch_device)
+    patch_device(model.encode_image)
+    patch_device(model.encode_text)
+
+    # patch dtype to float32 on CPU
+    if str(device) == "cpu":
+        float_holder = torch.jit.trace(lambda: torch.ones([]).float(), example_inputs=[])
+        float_input = list(float_holder.graph.findNode("aten::to").inputs())[1]
+        float_node = float_input.node()
+
+        def patch_float(module):
+            try:
+                graphs = [module.graph] if hasattr(module, "graph") else []
+            except RuntimeError:
+                graphs = []
+
+            if hasattr(module, "forward1"):
+                graphs.append(module.forward1.graph)
+
+            for graph in graphs:
+                for node in graph.findAllNodes("aten::to"):
+                    inputs = list(node.inputs())
+                    for i in [1, 2]:  # dtype can be the second or third argument to aten::to()
+                        if inputs[i].node()["value"] == 5:
+                            inputs[i].node().copyAttributes(float_node)
+
+        model.apply(patch_float)
+        patch_float(model.encode_image)
+        patch_float(model.encode_text)
+
+        model.float()
+
+    return model, _transform(model.input_resolution.item())
+
+
+def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]:
+    """
+    Returns the tokenized representation of given input string(s)
+
+    Parameters
+    ----------
+    texts : Union[str, List[str]]
+        An input string or a list of input strings to tokenize
+
+    context_length : int
+        The context length to use; all CLIP models use 77 as the context length
+
+    truncate: bool
+        Whether to truncate the text in case its encoding is longer than the context length
+
+    Returns
+    -------
+    A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length].
+    We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long.
+    """
+    if isinstance(texts, str):
+        texts = [texts]
+
+    sot_token = _tokenizer.encoder["<|startoftext|>"]
+    eot_token = _tokenizer.encoder["<|endoftext|>"]
+    all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts]
+    if version.parse(torch.__version__) < version.parse("1.8.0"):
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.long)
+    else:
+        result = torch.zeros(len(all_tokens), context_length, dtype=torch.int)
+
+    for i, tokens in enumerate(all_tokens):
+        if len(tokens) > context_length:
+            if truncate:
+                tokens = tokens[:context_length]
+                tokens[-1] = eot_token
+            else:
+                raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
+        result[i, :len(tokens)] = torch.tensor(tokens)
+
+    return result

models/dsp/CAM/clip/model.py

@@ -0,0 +1,517 @@
+from collections import OrderedDict
+from typing import Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+def upsample_pos_emb(emb, new_size):
+    # upsample the pretrained embedding for higher resolution
+    # emb size NxD
+    first = emb[:1, :]
+    emb = emb[1:, :]
+    N, D = emb.size(0), emb.size(1)
+    size = int(np.sqrt(N))
+    assert size * size == N
+    #new_size = size * self.upsample
+    emb = emb.permute(1, 0)
+    emb = emb.view(1, D, size, size).contiguous()
+    # emb = F.upsample(emb, size=new_size, mode='bilinear',)
+    emb = F.interpolate(emb, size=new_size, mode='bilinear', align_corners=False)
+    emb = emb.view(D, -1).contiguous()
+    emb = emb.permute(1, 0)
+    emb = torch.cat([first, emb], 0)
+    emb = nn.parameter.Parameter(emb.half())
+    return emb
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1):
+        super().__init__()
+
+        # all conv layers have stride 1. an avgpool is performed after the second convolution when stride > 1
+        self.conv1 = nn.Conv2d(inplanes, planes, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(planes)
+        self.relu1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(planes, planes, 3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.relu2 = nn.ReLU(inplace=True)
+
+        self.avgpool = nn.AvgPool2d(stride) if stride > 1 else nn.Identity()
+
+        self.conv3 = nn.Conv2d(planes, planes * self.expansion, 1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+        self.relu3 = nn.ReLU(inplace=True)
+
+        self.downsample = None
+        self.stride = stride
+
+        if stride > 1 or inplanes != planes * Bottleneck.expansion:
+            # downsampling layer is prepended with an avgpool, and the subsequent convolution has stride 1
+            self.downsample = nn.Sequential(OrderedDict([
+                ("-1", nn.AvgPool2d(stride)),
+                ("0", nn.Conv2d(inplanes, planes * self.expansion, 1, stride=1, bias=False)),
+                ("1", nn.BatchNorm2d(planes * self.expansion))
+            ]))
+
+    def forward(self, x: torch.Tensor):
+        identity = x
+
+        out = self.relu1(self.bn1(self.conv1(x)))
+        out = self.relu2(self.bn2(self.conv2(out)))
+        out = self.avgpool(out)
+        out = self.bn3(self.conv3(out))
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu3(out)
+        return out
+
+
+class AttentionPool2d(nn.Module):
+    def __init__(self, spacial_dim: int, embed_dim: int, num_heads: int, output_dim: int = None):
+        super().__init__()
+        self.positional_embedding = nn.Parameter(torch.randn(spacial_dim ** 2 + 1, embed_dim) / embed_dim ** 0.5)
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.c_proj = nn.Linear(embed_dim, output_dim or embed_dim)
+        self.num_heads = num_heads
+
+    def forward(self, x, H, W):
+        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
+        self.positional_embedding_new = upsample_pos_emb(self.positional_embedding, (H//32,W//32))
+        x = x + self.positional_embedding_new[:, None, :].to(x.dtype)  # (HW+1)NC
+        x, attn_weight = F.multi_head_attention_forward(
+            query=x, key=x, value=x,
+            embed_dim_to_check=x.shape[-1],
+            num_heads=self.num_heads,
+            q_proj_weight=self.q_proj.weight,
+            k_proj_weight=self.k_proj.weight,
+            v_proj_weight=self.v_proj.weight,
+            in_proj_weight=None,
+            in_proj_bias=torch.cat([self.q_proj.bias, self.k_proj.bias, self.v_proj.bias]),
+            bias_k=None,
+            bias_v=None,
+            add_zero_attn=False,
+            dropout_p=0,
+            out_proj_weight=self.c_proj.weight,
+            out_proj_bias=self.c_proj.bias,
+            use_separate_proj_weight=True,
+            training=self.training,
+            need_weights=False
+        )
+        return x[0]
+
+
+class ModifiedResNet(nn.Module):
+    """
+    A ResNet class that is similar to torchvision's but contains the following changes:
+    - There are now 3 "stem" convolutions as opposed to 1, with an average pool instead of a max pool.
+    - Performs anti-aliasing strided convolutions, where an avgpool is prepended to convolutions with stride > 1
+    - The final pooling layer is a QKV attention instead of an average pool
+    """
+
+    def __init__(self, layers, output_dim, heads, input_resolution=224, width=64):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_resolution = input_resolution
+
+        # the 3-layer stem
+        self.conv1 = nn.Conv2d(3, width // 2, kernel_size=3, stride=2, padding=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width // 2)
+        self.relu1 = nn.ReLU(inplace=True)
+        self.conv2 = nn.Conv2d(width // 2, width // 2, kernel_size=3, padding=1, bias=False)
+        self.bn2 = nn.BatchNorm2d(width // 2)
+        self.relu2 = nn.ReLU(inplace=True)
+        self.conv3 = nn.Conv2d(width // 2, width, kernel_size=3, padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(width)
+        self.relu3 = nn.ReLU(inplace=True)
+        self.avgpool = nn.AvgPool2d(2)
+
+        # residual layers
+        self._inplanes = width  # this is a *mutable* variable used during construction
+        self.layer1 = self._make_layer(width, layers[0])
+        self.layer2 = self._make_layer(width * 2, layers[1], stride=2)
+        self.layer3 = self._make_layer(width * 4, layers[2], stride=2)
+        self.layer4 = self._make_layer(width * 8, layers[3], stride=2)
+
+        embed_dim = width * 32  # the ResNet feature dimension
+        self.attnpool = AttentionPool2d(input_resolution // 32, embed_dim, heads, output_dim)
+
+    def _make_layer(self, planes, blocks, stride=1):
+        layers = [Bottleneck(self._inplanes, planes, stride)]
+
+        self._inplanes = planes * Bottleneck.expansion
+        for _ in range(1, blocks):
+            layers.append(Bottleneck(self._inplanes, planes))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x, H, W):
+        def stem(x):
+            x = self.relu1(self.bn1(self.conv1(x)))
+            x = self.relu2(self.bn2(self.conv2(x)))
+            x = self.relu3(self.bn3(self.conv3(x)))
+            x = self.avgpool(x)
+            return x
+
+        x = x.type(self.conv1.weight.dtype)
+        x = stem(x)
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)#(1,,2048, 7, 7)
+        x_pooled = self.attnpool(x, H, W)
+
+        return x_pooled
+
+
+class LayerNorm(nn.LayerNorm):
+    """Subclass torch's LayerNorm to handle fp16."""
+
+    def forward(self, x: torch.Tensor):
+        orig_type = x.dtype
+        ret = super().forward(x.type(torch.float32))
+        return ret.type(orig_type)
+
+
+class QuickGELU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class ResidualAttentionBlock(nn.Module):
+    def __init__(self, d_model: int, n_head: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+
+        self.attn = nn.MultiheadAttention(d_model, n_head)
+        self.ln_1 = LayerNorm(d_model)
+        self.mlp = nn.Sequential(OrderedDict([
+            ("c_fc", nn.Linear(d_model, d_model * 4)),
+            ("gelu", QuickGELU()),
+            ("c_proj", nn.Linear(d_model * 4, d_model))
+        ]))
+        self.ln_2 = LayerNorm(d_model)
+        self.attn_mask = attn_mask
+
+    def attention(self, x: torch.Tensor):
+        self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None
+        return self.attn(x, x, x, need_weights=True, attn_mask=self.attn_mask)#[0]
+
+    def forward(self, x: torch.Tensor):
+        attn_output, attn_weight = self.attention(self.ln_1(x))#(L,N,E)  (N,L,L)
+        x = x + attn_output
+        x = x + self.mlp(self.ln_2(x))
+        return x, attn_weight
+
+
+
+class Transformer(nn.Module):
+    def __init__(self, width: int, layers: int, heads: int, attn_mask: torch.Tensor = None):
+        super().__init__()
+        self.width = width
+        self.layers = layers
+        self.resblocks = nn.Sequential(*[ResidualAttentionBlock(width, heads, attn_mask) for _ in range(layers)])
+
+    def forward(self, x: torch.Tensor):
+        attn_weights = []
+        with torch.no_grad():
+            layers = self.layers if x.shape[0] == 77 else self.layers-1
+            for i in range(layers):
+                x, attn_weight = self.resblocks[i](x)
+                attn_weights.append(attn_weight)
+        '''
+        for i in range(self.layers-1, self.layers):
+            x, attn_weight = self.resblocks[i](x)
+            attn_weights.append(attn_weight)
+            #feature_map_list.append(x)
+        '''
+        return x, attn_weights
+
+
+class VisionTransformer(nn.Module):
+    def __init__(self, input_resolution: int, patch_size: int, width: int, layers: int, heads: int, output_dim: int):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.output_dim = output_dim
+        self.conv1 = nn.Conv2d(in_channels=3, out_channels=width, kernel_size=patch_size, stride=patch_size, bias=False)
+
+        scale = width ** -0.5
+        self.class_embedding = nn.Parameter(scale * torch.randn(width))
+        self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width))
+        self.ln_pre = LayerNorm(width)
+
+        self.transformer = Transformer(width, layers, heads)
+
+        self.ln_post = LayerNorm(width)
+        self.proj = nn.Parameter(scale * torch.randn(width, output_dim))
+        self.patch_size = patch_size
+
+    def forward(self, x: torch.Tensor, H, W):
+
+        self.positional_embedding_new = upsample_pos_emb(self.positional_embedding, (H//16,W//16))
+        x = self.conv1(x)  # shape = [*, width, grid, grid]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # shape = [*, width, grid ** 2]
+        x = x.permute(0, 2, 1)  # shape = [*, grid ** 2, width]
+        x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1)  # shape = [*, grid ** 2 + 1, width]
+        x = x + self.positional_embedding_new.to(x.dtype)
+        x = self.ln_pre(x)
+
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x, attn_weight = self.transformer(x)
+        '''
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.ln_post(x)
+        #x = x[:, 0, :]
+        #x = x[:,1:,:]
+        x = torch.mean(x[:,1:,:],dim=1)
+        #feature_map_list.append(x)
+
+        if self.proj is not None:
+            x = x @ self.proj
+        '''
+
+        return x, attn_weight#cls_attn
+
+
+class CLIP(nn.Module):
+    def __init__(self,
+                 embed_dim: int,
+                 # vision
+                 image_resolution: int,
+                 vision_layers: Union[Tuple[int, int, int, int], int],
+                 vision_width: int,
+                 vision_patch_size: int,
+                 # text
+                 context_length: int,
+                 vocab_size: int,
+                 transformer_width: int,
+                 transformer_heads: int,
+                 transformer_layers: int
+                 ):
+        super().__init__()
+
+        self.context_length = context_length
+
+        if isinstance(vision_layers, (tuple, list)):
+            vision_heads = vision_width * 32 // 64
+            self.visual = ModifiedResNet(
+                layers=vision_layers,
+                output_dim=embed_dim,
+                heads=vision_heads,
+                input_resolution=image_resolution,
+                width=vision_width
+            )
+        else:
+            vision_heads = vision_width // 64
+            self.visual = VisionTransformer(
+                input_resolution=image_resolution,
+                patch_size=vision_patch_size,
+                width=vision_width,
+                layers=vision_layers,
+                heads=vision_heads,
+                output_dim=embed_dim
+            )
+
+        self.transformer = Transformer(
+            width=transformer_width,
+            layers=transformer_layers,
+            heads=transformer_heads,
+            attn_mask=self.build_attention_mask()
+        )
+
+        self.vocab_size = vocab_size
+        self.token_embedding = nn.Embedding(vocab_size, transformer_width)
+        self.positional_embedding = nn.Parameter(torch.empty(self.context_length, transformer_width))
+        self.ln_final = LayerNorm(transformer_width)
+
+        self.text_projection = nn.Parameter(torch.empty(transformer_width, embed_dim))
+        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
+
+        self.initialize_parameters()
+
+    def initialize_parameters(self):
+        nn.init.normal_(self.token_embedding.weight, std=0.02)
+        nn.init.normal_(self.positional_embedding, std=0.01)
+
+        if isinstance(self.visual, ModifiedResNet):
+            if self.visual.attnpool is not None:
+                std = self.visual.attnpool.c_proj.in_features ** -0.5
+                nn.init.normal_(self.visual.attnpool.q_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.k_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.v_proj.weight, std=std)
+                nn.init.normal_(self.visual.attnpool.c_proj.weight, std=std)
+
+            for resnet_block in [self.visual.layer1, self.visual.layer2, self.visual.layer3, self.visual.layer4]:
+                for name, param in resnet_block.named_parameters():
+                    if name.endswith("bn3.weight"):
+                        nn.init.zeros_(param)
+
+        proj_std = (self.transformer.width ** -0.5) * ((2 * self.transformer.layers) ** -0.5)
+        attn_std = self.transformer.width ** -0.5
+        fc_std = (2 * self.transformer.width) ** -0.5
+        for block in self.transformer.resblocks:
+            nn.init.normal_(block.attn.in_proj_weight, std=attn_std)
+            nn.init.normal_(block.attn.out_proj.weight, std=proj_std)
+            nn.init.normal_(block.mlp.c_fc.weight, std=fc_std)
+            nn.init.normal_(block.mlp.c_proj.weight, std=proj_std)
+
+        if self.text_projection is not None:
+            nn.init.normal_(self.text_projection, std=self.transformer.width ** -0.5)
+
+    def build_attention_mask(self):
+        # lazily create causal attention mask, with full attention between the vision tokens
+        # pytorch uses additive attention mask; fill with -inf
+        mask = torch.empty(self.context_length, self.context_length)
+        mask.fill_(float("-inf"))
+        mask.triu_(1)  # zero out the lower diagonal
+        return mask
+
+    @property
+    def dtype(self):
+        return self.visual.conv1.weight.dtype
+
+    def encode_image(self, image, H, W):
+        return self.visual(image.type(self.dtype), H, W)
+
+    def encode_text(self, text):
+        x = self.token_embedding(text).type(self.dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.type(self.dtype)
+        x = x.permute(1, 0, 2)  # NLD -> LND
+        x, attn_weight = self.transformer(x)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+        x = self.ln_final(x).type(self.dtype)
+
+        # x.shape = [batch_size, n_ctx, transformer.width]
+        # take features from the eot embedding (eot_token is the highest number in each sequence)
+        x = x[torch.arange(x.shape[0]), text.argmax(dim=-1)] @ self.text_projection
+
+        return x
+
+    def forward_last_layer(self, image_features, text_features):
+        x, attn_weight = self.visual.transformer.resblocks[self.visual.transformer.layers-1](image_features)
+        x = x.permute(1, 0, 2)  # LND -> NLD
+
+        x = self.visual.ln_post(x)
+        x = torch.mean(x[:, 1:, :], dim=1)
+
+        if self.visual.proj is not None:
+            x = x @ self.visual.proj
+
+        image_features = x
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        logits_per_image = logits_per_image.softmax(dim=-1)
+
+        return logits_per_image, attn_weight
+
+
+
+
+    def forward(self, image, text):
+        image_features, feature_map, cls_attn = self.encode_image(image)
+        with torch.no_grad():
+            text_features = self.encode_text(text)
+
+        # normalized features
+        image_features = image_features / image_features.norm(dim=1, keepdim=True)
+        text_features = text_features / text_features.norm(dim=1, keepdim=True)
+
+        # cosine similarity as logits
+        logit_scale = self.logit_scale.exp()
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        #logits_per_text = logits_per_image.t()
+
+        # shape = [global_batch_size, global_batch_size]
+        return logits_per_image, logits_per_text
+
+
+def convert_weights(model: nn.Module):
+    """Convert applicable model parameters to fp16"""
+
+    def _convert_weights_to_fp16(l):
+        if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Linear)):
+            l.weight.data = l.weight.data.half()
+            if l.bias is not None:
+                l.bias.data = l.bias.data.half()
+
+        if isinstance(l, nn.MultiheadAttention):
+            for attr in [*[f"{s}_proj_weight" for s in ["in", "q", "k", "v"]], "in_proj_bias", "bias_k", "bias_v"]:
+                tensor = getattr(l, attr)
+                if tensor is not None:
+                    tensor.data = tensor.data.half()
+
+        for name in ["text_projection", "proj"]:
+            if hasattr(l, name):
+                attr = getattr(l, name)
+                if attr is not None:
+                    attr.data = attr.data.half()
+
+    model.apply(_convert_weights_to_fp16)
+
+
+def build_model(state_dict: dict):
+    vit = "visual.proj" in state_dict
+    '''
+    inv_freq = 1. / (10000 ** (torch.arange(0, 2048, 2, dtype=torch.float) / 2048))
+    position = torch.arange(50, dtype=torch.float)
+    sinusoid_inp = torch.einsum('i,j -> ij', position, inv_freq)
+    embeddings = torch.cat((sinusoid_inp.sin(), sinusoid_inp.cos()), dim=-1) / 2048 ** 0.5
+    state_dict["visual.attnpool.positional_embedding"] = embeddings
+    '''
+    #state_dict["visual.positional_embedding"] = upsample_pos_emb(state_dict["visual.positional_embedding"], 28)
+
+
+    if vit:
+        vision_width = state_dict["visual.conv1.weight"].shape[0]
+        vision_layers = len([k for k in state_dict.keys() if k.startswith("visual.") and k.endswith(".attn.in_proj_weight")])
+        vision_patch_size = state_dict["visual.conv1.weight"].shape[-1]
+        grid_size = round((state_dict["visual.positional_embedding"].shape[0] - 1) ** 0.5)
+        image_resolution = vision_patch_size * grid_size
+    else:
+        counts: list = [len(set(k.split(".")[2] for k in state_dict if k.startswith(f"visual.layer{b}"))) for b in [1, 2, 3, 4]]
+        vision_layers = tuple(counts)
+        vision_width = state_dict["visual.layer1.0.conv1.weight"].shape[0]
+        output_width = round((state_dict["visual.attnpool.positional_embedding"].shape[0] - 1) ** 0.5)
+        vision_patch_size = None
+        assert output_width ** 2 + 1 == state_dict["visual.attnpool.positional_embedding"].shape[0]
+        image_resolution = output_width * 32
+
+    embed_dim = state_dict["text_projection"].shape[1]
+    context_length = state_dict["positional_embedding"].shape[0]
+    vocab_size = state_dict["token_embedding.weight"].shape[0]
+    transformer_width = state_dict["ln_final.weight"].shape[0]
+    transformer_heads = transformer_width // 64
+    transformer_layers = len(set(k.split(".")[2] for k in state_dict if k.startswith(f"transformer.resblocks")))
+
+    model = CLIP(
+        embed_dim,
+        image_resolution, vision_layers, vision_width, vision_patch_size,
+        context_length, vocab_size, transformer_width, transformer_heads, transformer_layers
+    )
+
+    for key in ["input_resolution", "context_length", "vocab_size"]:
+        if key in state_dict:
+            del state_dict[key]
+
+
+
+    convert_weights(model)
+    model.load_state_dict(state_dict)
+    return model.eval()

models/dsp/CAM/clip/simple_tokenizer.py

@@ -0,0 +1,132 @@
+import gzip
+import html
+import os
+from functools import lru_cache
+
+import ftfy
+import regex as re
+
+
+@lru_cache()
+def default_bpe():
+    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "bpe_simple_vocab_16e6.txt.gz")
+
+
+@lru_cache()
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a corresponding list of unicode strings.
+    The reversible bpe codes work on unicode strings.
+    This means you need a large # of unicode characters in your vocab if you want to avoid UNKs.
+    When you're at something like a 10B token dataset you end up needing around 5K for decent coverage.
+    This is a signficant percentage of your normal, say, 32K bpe vocab.
+    To avoid that, we want lookup tables between utf-8 bytes and unicode strings.
+    And avoids mapping to whitespace/control characters the bpe code barfs on.
+    """
+    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8+n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+def get_pairs(word):
+    """Return set of symbol pairs in a word.
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+def basic_clean(text):
+    text = ftfy.fix_text(text)
+    text = html.unescape(html.unescape(text))
+    return text.strip()
+
+
+def whitespace_clean(text):
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    return text
+
+
+class SimpleTokenizer(object):
+    def __init__(self, bpe_path: str = default_bpe()):
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        merges = gzip.open(bpe_path).read().decode("utf-8").split('\n')
+        merges = merges[1:49152-256-2+1]
+        merges = [tuple(merge.split()) for merge in merges]
+        vocab = list(bytes_to_unicode().values())
+        vocab = vocab + [v+'</w>' for v in vocab]
+        for merge in merges:
+            vocab.append(''.join(merge))
+        vocab.extend(['<|startoftext|>', '<|endoftext|>'])
+        self.encoder = dict(zip(vocab, range(len(vocab))))
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.bpe_ranks = dict(zip(merges, range(len(merges))))
+        self.cache = {'<|startoftext|>': '<|startoftext|>', '<|endoftext|>': '<|endoftext|>'}
+        self.pat = re.compile(r"""<\|startoftext\|>|<\|endoftext\|>|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""", re.IGNORECASE)
+
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token[:-1]) + ( token[-1] + '</w>',)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token+'</w>'
+
+        while True:
+            bigram = min(pairs, key = lambda pair: self.bpe_ranks.get(pair, float('inf')))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                    new_word.extend(word[i:j])
+                    i = j
+                except:
+                    new_word.extend(word[i:])
+                    break
+
+                if word[i] == first and i < len(word)-1 and word[i+1] == second:
+                    new_word.append(first+second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = ' '.join(word)
+        self.cache[token] = word
+        return word
+
+    def encode(self, text):
+        bpe_tokens = []
+        text = whitespace_clean(basic_clean(text)).lower()
+        for token in re.findall(self.pat, text):
+            token = ''.join(self.byte_encoder[b] for b in token.encode('utf-8'))
+            bpe_tokens.extend(self.encoder[bpe_token] for bpe_token in self.bpe(token).split(' '))
+        return bpe_tokens
+
+    def decode(self, tokens):
+        text = ''.join([self.decoder[token] for token in tokens])
+        text = bytearray([self.byte_decoder[c] for c in text]).decode('utf-8', errors="replace").replace('</w>', ' ')
+        return text

models/dsp/CAM/pytorch_grad_cam/__init__.py

@@ -0,0 +1,14 @@
+from .grad_cam import GradCAM
+from .ablation_layer import AblationLayer, AblationLayerVit, AblationLayerFasterRCNN
+from .ablation_cam import AblationCAM
+from .xgrad_cam import XGradCAM
+from .grad_cam_plusplus import GradCAMPlusPlus
+from .score_cam import ScoreCAM
+from .layer_cam import LayerCAM
+from .eigen_cam import EigenCAM
+from .eigen_grad_cam import EigenGradCAM
+from .fullgrad_cam import FullGrad
+from .guided_backprop import GuidedBackpropReLUModel
+from .activations_and_gradients import ActivationsAndGradients
+from .utils import model_targets
+from .utils import reshape_transforms

models/dsp/CAM/pytorch_grad_cam/ablation_cam.py

@@ -0,0 +1,134 @@
+import numpy as np
+import torch
+import tqdm
+from typing import Callable, List
+from .base_cam import BaseCAM
+from .utils.find_layers import replace_layer_recursive
+from .ablation_layer import AblationLayer
+
+
+""" Implementation of AblationCAM
+https://openaccess.thecvf.com/content_WACV_2020/papers/Desai_Ablation-CAM_Visual_Explanations_for_Deep_Convolutional_Network_via_Gradient-free_Localization_WACV_2020_paper.pdf
+
+Ablate individual activations, and then measure the drop in the target score.
+
+In the current implementation, the target layer activations is cached, so it won't be re-computed.
+However layers before it, if any, will not be cached.
+This means that if the target layer is a large block, for example model.featuers (in vgg), there will
+be a large save in run time.
+
+Since we have to go over many channels and ablate them, and every channel ablation requires a forward pass,
+it would be nice if we could avoid doing that for channels that won't contribute anwyay, making it much faster.
+The parameter ratio_channels_to_ablate controls how many channels should be ablated, using an experimental method
+(to be improved). The default 1.0 value means that all channels will be ablated.
+"""
+
+
+class AblationCAM(BaseCAM):
+    def __init__(self,
+                 model: torch.nn.Module,
+                 target_layers: List[torch.nn.Module],
+                 use_cuda: bool = False,
+                 reshape_transform: Callable = None,
+                 ablation_layer: torch.nn.Module = AblationLayer(),
+                 batch_size: int = 32,
+                 ratio_channels_to_ablate: float = 1.0) -> None:
+
+        super(AblationCAM, self).__init__(model,
+                                          target_layers,
+                                          use_cuda,
+                                          reshape_transform,
+                                          uses_gradients=False)
+        self.batch_size = batch_size
+        self.ablation_layer = ablation_layer
+        self.ratio_channels_to_ablate = ratio_channels_to_ablate
+
+    def save_activation(self, module, input, output) -> None:
+        """ Helper function to save the raw activations from the target layer """
+        self.activations = output
+
+    def assemble_ablation_scores(self,
+                                 new_scores: list,
+                                 original_score: float ,
+                                 ablated_channels: np.ndarray,
+                                 number_of_channels: int) -> np.ndarray:
+        """ Take the value from the channels that were ablated,
+            and just set the original score for the channels that were skipped """
+
+        index = 0
+        result = []
+        sorted_indices = np.argsort(ablated_channels)
+        ablated_channels = ablated_channels[sorted_indices]
+        new_scores = np.float32(new_scores)[sorted_indices]
+
+        for i in range(number_of_channels):
+            if index < len(ablated_channels) and ablated_channels[index] == i:
+                weight = new_scores[index]
+                index = index + 1
+            else:
+                weight = original_score
+            result.append(weight)
+
+        return result
+
+    def get_cam_weights(self,
+                        input_tensor: torch.Tensor,
+                        target_layer: torch.nn.Module,
+                        targets: List[Callable],
+                        activations: torch.Tensor,
+                        grads: torch.Tensor) -> np.ndarray:
+        
+        # Do a forward pass, compute the target scores, and cache the activations
+        handle = target_layer.register_forward_hook(self.save_activation)
+        with torch.no_grad():
+            outputs = self.model(input_tensor)
+            handle.remove()
+            original_scores = np.float32([target(output).cpu().item() for target, output in zip(targets, outputs)])
+
+        # Replace the layer with the ablation layer.
+        # When we finish, we will replace it back, so the original model is unchanged.
+        ablation_layer = self.ablation_layer
+        replace_layer_recursive(self.model, target_layer, ablation_layer)
+
+        number_of_channels = activations.shape[1]
+        weights = []
+        # This is a "gradient free" method, so we don't need gradients here.
+        with torch.no_grad():
+            # Loop over each of the batch images and ablate activations for it.
+            for batch_index, (target, tensor) in enumerate(zip(targets, input_tensor)):
+                new_scores = []
+                batch_tensor = tensor.repeat(self.batch_size, 1, 1, 1)
+
+                # Check which channels should be ablated. Normally this will be all channels,
+                # But we can also try to speed this up by using a low ratio_channels_to_ablate.
+                channels_to_ablate = ablation_layer.activations_to_be_ablated(activations[batch_index, :],
+                                                                              self.ratio_channels_to_ablate)
+                number_channels_to_ablate = len(channels_to_ablate)
+
+                for i in tqdm.tqdm(range(0, number_channels_to_ablate, self.batch_size)):
+                    if i + self.batch_size > number_channels_to_ablate:
+                        batch_tensor = batch_tensor[:(number_channels_to_ablate - i)]
+
+                    # Change the state of the ablation layer so it ablates the next channels.
+                    # TBD: Move this into the ablation layer forward pass.
+                    ablation_layer.set_next_batch(input_batch_index=batch_index,
+                                                  activations=self.activations,
+                                                  num_channels_to_ablate=batch_tensor.size(0))
+                    score = [target(o).cpu().item() for o in self.model(batch_tensor)]
+                    new_scores.extend(score)
+                    ablation_layer.indices = ablation_layer.indices[batch_tensor.size(0):]
+
+                new_scores = self.assemble_ablation_scores(new_scores,
+                                                           original_scores[batch_index],
+                                                           channels_to_ablate,
+                                                           number_of_channels)
+                weights.extend(new_scores)
+
+        weights = np.float32(weights)
+        weights = weights.reshape(activations.shape[:2])
+        original_scores = original_scores[:, None]
+        weights = (original_scores - weights) / original_scores
+
+        # Replace the model back to the original state
+        replace_layer_recursive(self.model, ablation_layer, target_layer)
+        return weights

models/dsp/CAM/pytorch_grad_cam/ablation_cam_multilayer.py

@@ -0,0 +1,136 @@
+import cv2
+import numpy as np
+import torch
+import tqdm
+from .base_cam import BaseCAM
+
+
+class AblationLayer(torch.nn.Module):
+    def __init__(self, layer, reshape_transform, indices):
+        super(AblationLayer, self).__init__()
+
+        self.layer = layer
+        self.reshape_transform = reshape_transform
+        # The channels to zero out:
+        self.indices = indices
+
+    def forward(self, x):
+        self.__call__(x)
+
+    def __call__(self, x):
+        output = self.layer(x)
+
+        # Hack to work with ViT,
+        # Since the activation channels are last and not first like in CNNs
+        # Probably should remove it?
+        if self.reshape_transform is not None:
+            output = output.transpose(1, 2)
+
+        for i in range(output.size(0)):
+
+            # Commonly the minimum activation will be 0,
+            # And then it makes sense to zero it out.
+            # However depending on the architecture,
+            # If the values can be negative, we use very negative values
+            # to perform the ablation, deviating from the paper.
+            if torch.min(output) == 0:
+                output[i, self.indices[i], :] = 0
+            else:
+                ABLATION_VALUE = 1e5
+                output[i, self.indices[i], :] = torch.min(
+                    output) - ABLATION_VALUE
+
+        if self.reshape_transform is not None:
+            output = output.transpose(2, 1)
+
+        return output
+
+
+def replace_layer_recursive(model, old_layer, new_layer):
+    for name, layer in model._modules.items():
+        if layer == old_layer:
+            model._modules[name] = new_layer
+            return True
+        elif replace_layer_recursive(layer, old_layer, new_layer):
+            return True
+    return False
+
+
+class AblationCAM(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        super(AblationCAM, self).__init__(model, target_layers, use_cuda,
+                                          reshape_transform)
+
+        if len(target_layers) > 1:
+            print(
+                "Warning. You are usign Ablation CAM with more than 1 layers. "
+                "This is supported only if all layers have the same output shape")
+
+    def set_ablation_layers(self):
+        self.ablation_layers = []
+        for target_layer in self.target_layers:
+            ablation_layer = AblationLayer(target_layer,
+                                           self.reshape_transform, indices=[])
+            self.ablation_layers.append(ablation_layer)
+            replace_layer_recursive(self.model, target_layer, ablation_layer)
+
+    def unset_ablation_layers(self):
+        # replace the model back to the original state
+        for ablation_layer, target_layer in zip(
+                self.ablation_layers, self.target_layers):
+            replace_layer_recursive(self.model, ablation_layer, target_layer)
+
+    def set_ablation_layer_batch_indices(self, indices):
+        for ablation_layer in self.ablation_layers:
+            ablation_layer.indices = indices
+
+    def trim_ablation_layer_batch_indices(self, keep):
+        for ablation_layer in self.ablation_layers:
+            ablation_layer.indices = ablation_layer.indices[:keep]
+
+    def get_cam_weights(self,
+                        input_tensor,
+                        target_category,
+                        activations,
+                        grads):
+        with torch.no_grad():
+            outputs = self.model(input_tensor).cpu().numpy()
+            original_scores = []
+            for i in range(input_tensor.size(0)):
+                original_scores.append(outputs[i, target_category[i]])
+        original_scores = np.float32(original_scores)
+
+        self.set_ablation_layers()
+
+        if hasattr(self, "batch_size"):
+            BATCH_SIZE = self.batch_size
+        else:
+            BATCH_SIZE = 32
+
+        number_of_channels = activations.shape[1]
+        weights = []
+
+        with torch.no_grad():
+            # Iterate over the input batch
+            for tensor, category in zip(input_tensor, target_category):
+                batch_tensor = tensor.repeat(BATCH_SIZE, 1, 1, 1)
+                for i in tqdm.tqdm(range(0, number_of_channels, BATCH_SIZE)):
+                    self.set_ablation_layer_batch_indices(
+                        list(range(i, i + BATCH_SIZE)))
+
+                    if i + BATCH_SIZE > number_of_channels:
+                        keep = number_of_channels - i
+                        batch_tensor = batch_tensor[:keep]
+                        self.trim_ablation_layer_batch_indices(self, keep)
+                    score = self.model(batch_tensor)[:, category].cpu().numpy()
+                    weights.extend(score)
+
+        weights = np.float32(weights)
+        weights = weights.reshape(activations.shape[:2])
+        original_scores = original_scores[:, None]
+        weights = (original_scores - weights) / original_scores
+
+        # replace the model back to the original state
+        self.unset_ablation_layers()
+        return weights

models/dsp/CAM/pytorch_grad_cam/ablation_layer.py

@@ -0,0 +1,131 @@
+import torch
+from collections import OrderedDict
+import numpy as np
+from .utils.svd_on_activations import get_2d_projection
+
+
+class AblationLayer(torch.nn.Module):
+    def __init__(self):
+        super(AblationLayer, self).__init__()
+
+    def objectiveness_mask_from_svd(self, activations, threshold=0.01):
+        """ Experimental method to get a binary mask to compare if the activation is worth ablating.
+            The idea is to apply the EigenCAM method by doing PCA on the activations.
+            Then we create a binary mask by comparing to a low threshold.
+            Areas that are masked out, are probably not interesting anyway.
+        """
+
+        projection = get_2d_projection(activations[None, :])[0, :]
+        projection = np.abs(projection)
+        projection = projection - projection.min()
+        projection = projection / projection.max()
+        projection = projection > threshold
+        return projection
+
+    def activations_to_be_ablated(self, activations, ratio_channels_to_ablate=1.0):
+        """ Experimental method to get a binary mask to compare if the activation is worth ablating.
+            Create a binary CAM mask with objectiveness_mask_from_svd.
+            Score each Activation channel, by seeing how much of its values are inside the mask.
+            Then keep the top channels.
+
+        """
+        if ratio_channels_to_ablate == 1.0:
+            self.indices = np.int32(range(activations.shape[0]))
+            return self.indices
+
+        projection = self.objectiveness_mask_from_svd(activations)
+
+        scores = []
+        for channel in activations:
+            normalized = np.abs(channel)
+            normalized = normalized - normalized.min()
+            normalized = normalized / np.max(normalized)
+            score = (projection*normalized).sum() / normalized.sum()
+            scores.append(score)
+        scores = np.float32(scores)
+
+        indices = list(np.argsort(scores))
+        high_score_indices = indices[::-1][: int(len(indices) * ratio_channels_to_ablate)]
+        low_score_indices = indices[: int(len(indices) * ratio_channels_to_ablate)]
+        self.indices = np.int32(high_score_indices + low_score_indices)
+        return self.indices
+
+    def set_next_batch(self, input_batch_index, activations, num_channels_to_ablate):
+        """ This creates the next batch of activations from the layer.
+            Just take corresponding batch member from activations, and repeat it num_channels_to_ablate times.
+        """
+        self.activations = activations[input_batch_index, :, :, :].clone().unsqueeze(0).repeat(num_channels_to_ablate, 1, 1, 1)
+
+    def __call__(self, x):
+        output = self.activations
+        for i in range(output.size(0)):
+            # Commonly the minimum activation will be 0,
+            # And then it makes sense to zero it out.
+            # However depending on the architecture,
+            # If the values can be negative, we use very negative values
+            # to perform the ablation, deviating from the paper.
+            if torch.min(output) == 0:
+                output[i, self.indices[i], :] = 0
+            else:
+                ABLATION_VALUE = 1e7
+                output[i, self.indices[i], :] = torch.min(
+                    output) - ABLATION_VALUE
+
+        return output
+
+
+class AblationLayerVit(AblationLayer):
+    def __init__(self):
+        super(AblationLayerVit, self).__init__()
+
+    def __call__(self, x):
+        output = self.activations
+        output = output.transpose(1, 2)
+        for i in range(output.size(0)):
+
+            # Commonly the minimum activation will be 0,
+            # And then it makes sense to zero it out.
+            # However depending on the architecture,
+            # If the values can be negative, we use very negative values
+            # to perform the ablation, deviating from the paper.
+            if torch.min(output) == 0:
+                output[i, self.indices[i], :] = 0
+            else:
+                ABLATION_VALUE = 1e7
+                output[i, self.indices[i], :] = torch.min(
+                    output) - ABLATION_VALUE
+
+        output = output.transpose(2, 1)
+
+        return output
+
+    def set_next_batch(self, input_batch_index, activations, num_channels_to_ablate):
+        """ This creates the next batch of activations from the layer.
+            Just take corresponding batch member from activations, and repeat it num_channels_to_ablate times.
+        """
+        self.activations = activations[input_batch_index, :, :].clone().unsqueeze(0).repeat(num_channels_to_ablate, 1, 1)
+
+
+
+class AblationLayerFasterRCNN(AblationLayer):
+    def __init__(self):
+        super(AblationLayerFasterRCNN, self).__init__()
+
+    def set_next_batch(self, input_batch_index, activations, num_channels_to_ablate):
+        """ Extract the next batch member from activations,
+            and repeat it num_channels_to_ablate times.
+        """
+        self.activations = OrderedDict()
+        for key, value in activations.items():
+            fpn_activation = value[input_batch_index, :, :, :].clone().unsqueeze(0)
+            self.activations[key] = fpn_activation.repeat(num_channels_to_ablate, 1, 1, 1)
+
+    def __call__(self, x):
+        result = self.activations
+        layers = {0: '0', 1: '1', 2: '2', 3: '3', 4: 'pool'}
+        num_channels_to_ablate = result['pool'].size(0)
+        for i in range(num_channels_to_ablate):
+            pyramid_layer = int(self.indices[i]/256)
+            index_in_pyramid_layer = int(self.indices[i] % 256)
+            result[layers[pyramid_layer]][i, index_in_pyramid_layer, :, :] = -1000
+        return result

models/dsp/CAM/pytorch_grad_cam/activations_and_gradients.py

@@ -0,0 +1,55 @@
+class ActivationsAndGradients:
+    """ Class for extracting activations and
+    registering gradients from targetted intermediate layers """
+
+    def __init__(self, model, target_layers, reshape_transform):
+        self.model = model
+        self.gradients = []
+        self.activations = []
+        self.reshape_transform = reshape_transform
+        self.handles = []
+        for target_layer in target_layers:
+            self.handles.append(
+                target_layer.register_forward_hook(self.save_activation))
+            # Because of https://github.com/pytorch/pytorch/issues/61519,
+            # we don't use backward hook to record gradients.
+            self.handles.append(
+                target_layer.register_forward_hook(self.save_gradient))
+
+    def save_activation(self, module, input, output):
+        activation = output
+
+        if self.reshape_transform is not None:
+            activation = self.reshape_transform(activation, self.height, self.width)
+        # self.activations.append(activation.cpu().detach())
+        # self.activations.append(activation.detach()) # original detach
+        self.activations.append(activation)
+
+    def save_gradient(self, module, input, output):
+        if not hasattr(output, "requires_grad") or not output.requires_grad:
+            # You can only register hooks on tensor requires grad.
+            return
+
+        # Gradients are computed in reverse order
+        def _store_grad(grad):
+            if self.reshape_transform is not None:
+                grad = self.reshape_transform(grad, self.height, self.width)
+            # self.gradients = [grad.cpu().detach()] + self.gradients
+            # self.gradients = [grad.detach()] + self.gradients # original detach
+            self.gradients = [grad] + self.gradients
+
+        output.register_hook(_store_grad)
+
+    def __call__(self, x, H, W):
+        self.height = H // 16
+        self.width = W // 16
+        self.gradients = []
+        self.activations = []
+        if isinstance(x, list):
+            return self.model.forward_last_layer(x[0], x[1])
+        else:
+            return self.model(x)
+
+    def release(self):
+        for handle in self.handles:
+            handle.remove()

models/dsp/CAM/pytorch_grad_cam/base_cam.py

@@ -0,0 +1,227 @@
+import numpy as np
+import torch
+import ttach as tta
+from typing import Callable, List, Tuple
+from .activations_and_gradients import ActivationsAndGradients
+from .utils.svd_on_activations import get_2d_projection
+from .utils.image import scale_cam_image, scale_cam_image_torch
+from .utils.model_targets import ClassifierOutputTarget
+
+
+class BaseCAM:
+    def __init__(self,
+                 model: torch.nn.Module,
+                 target_layers: List[torch.nn.Module],
+                 use_cuda: bool = False,
+                 reshape_transform: Callable = None,
+                 compute_input_gradient: bool = False,
+                 uses_gradients: bool = True) -> None:
+        self.model = model.eval()
+        self.target_layers = target_layers
+        self.cuda = use_cuda
+        # if self.cuda:
+        #     self.model = model.cuda()
+        self.reshape_transform = reshape_transform
+        self.compute_input_gradient = compute_input_gradient
+        self.uses_gradients = uses_gradients
+        self.activations_and_grads = ActivationsAndGradients(
+            self.model, target_layers, reshape_transform)
+
+    """ Get a vector of weights for every channel in the target layer.
+        Methods that return weights channels,
+        will typically need to only implement this function. """
+    
+    def set_device(self, device):
+        self.device = device
+
+    def get_cam_weights(self,
+                        input_tensor: torch.Tensor,
+                        target_layers: List[torch.nn.Module],
+                        targets: List[torch.nn.Module],
+                        activations: torch.Tensor,
+                        grads: torch.Tensor) -> np.ndarray:
+        raise Exception("Not Implemented")
+
+    def get_cam_image(self,
+                      input_tensor: torch.Tensor,
+                      target_layer: torch.nn.Module,
+                      targets: List[torch.nn.Module],
+                      activations: torch.Tensor,
+                      grads: torch.Tensor,
+                      eigen_smooth: bool = False) -> np.ndarray:
+
+        weights = self.get_cam_weights(input_tensor,
+                                       target_layer,
+                                       targets,
+                                       activations,
+                                       grads)
+        weighted_activations = weights[:, :, None, None] * activations
+        if eigen_smooth:
+            cam = get_2d_projection(weighted_activations)
+        else:
+            cam = weighted_activations.sum(axis=1)
+        return cam
+
+    def forward(self,
+                input_tensor: torch.Tensor,
+                targets: List[torch.nn.Module],
+                target_size,
+                eigen_smooth: bool = False) -> np.ndarray:
+
+        if self.cuda:
+            if isinstance(input_tensor, list):
+                input_tensor = [t.to(self.device) if isinstance(t, torch.Tensor) else t for t in input_tensor]
+            else:
+                input_tensor = input_tensor.to(self.device)
+
+        if self.compute_input_gradient:
+            input_tensor = torch.autograd.Variable(input_tensor,
+                                                   requires_grad=True)
+
+        W,H = self.get_target_width_height(input_tensor)
+        outputs = self.activations_and_grads(input_tensor,H,W)
+        if targets is None:
+            if isinstance(input_tensor, list):
+                target_categories = np.argmax(outputs[0].cpu().data.numpy(), axis=-1)
+            else:
+                target_categories = np.argmax(outputs.cpu().data.numpy(), axis=-1)
+            targets = [ClassifierOutputTarget(category) for category in target_categories]
+
+        if self.uses_gradients:
+            self.model.zero_grad()
+            if isinstance(input_tensor, list):
+                loss = sum([target(output[0]) for target, output in zip(targets, outputs)])
+            else:
+                loss = sum([target(output) for target, output in zip(targets, outputs)])
+            loss.backward(retain_graph=True)
+
+        # In most of the saliency attribution papers, the saliency is
+        # computed with a single target layer.
+        # Commonly it is the last convolutional layer.
+        # Here we support passing a list with multiple target layers.
+        # It will compute the saliency image for every image,
+        # and then aggregate them (with a default mean aggregation).
+        # This gives you more flexibility in case you just want to
+        # use all conv layers for example, all Batchnorm layers,
+        # or something else.
+        cam_per_layer = self.compute_cam_per_layer(input_tensor,
+                                                   targets,
+                                                   target_size,
+                                                   eigen_smooth)
+        if isinstance(input_tensor, list):
+            return self.aggregate_multi_layers_torch(cam_per_layer), outputs[0], outputs[1]
+        else:
+            return self.aggregate_multi_layers(cam_per_layer), outputs
+
+    def get_target_width_height(self,
+                                input_tensor: torch.Tensor) -> Tuple[int, int]:
+        if isinstance(input_tensor, list):
+            width, height = input_tensor[-1], input_tensor[-2]
+        return width, height
+
+    def compute_cam_per_layer(
+            self,
+            input_tensor: torch.Tensor,
+            targets: List[torch.nn.Module],
+            target_size,
+            eigen_smooth: bool) -> np.ndarray:
+        activations_list = [a#.cpu().data.numpy()
+                            for a in self.activations_and_grads.activations]
+        grads_list = [g#.cpu().data.numpy()
+                      for g in self.activations_and_grads.gradients]
+
+        cam_per_target_layer = []
+        # Loop over the saliency image from every layer
+        for i in range(len(self.target_layers)):
+            target_layer = self.target_layers[i]
+            layer_activations = None
+            layer_grads = None
+            if i < len(activations_list):
+                layer_activations = activations_list[i]
+            if i < len(grads_list):
+                layer_grads = grads_list[i]
+
+            cam = self.get_cam_image(input_tensor,
+                                     target_layer,
+                                     targets,
+                                     layer_activations,
+                                     layer_grads,
+                                     eigen_smooth)
+            cam = torch.clamp(cam, min=0).float()
+            scaled = scale_cam_image_torch(cam, target_size)
+            # cam = np.maximum(cam, 0).astype(np.float32)#float16->32
+            # scaled = scale_cam_image(cam, target_size)
+            cam_per_target_layer.append(scaled[:, None, :])
+
+        return cam_per_target_layer
+
+    def aggregate_multi_layers(self, cam_per_target_layer: np.ndarray) -> np.ndarray:
+        cam_per_target_layer = np.concatenate(cam_per_target_layer, axis=1)
+        cam_per_target_layer = np.maximum(cam_per_target_layer, 0)
+        result = np.mean(cam_per_target_layer, axis=1)
+        return scale_cam_image(result)
+    
+    def aggregate_multi_layers_torch(self, cam_per_target_layer):
+        cam_per_target_layer = torch.cat(cam_per_target_layer, dim=1)
+        cam_per_target_layer = torch.clamp(cam_per_target_layer, min=0)
+        result = cam_per_target_layer.mean(dim=1, keepdim=True)
+        return scale_cam_image_torch(result)
+
+    def forward_augmentation_smoothing(self,
+                                       input_tensor: torch.Tensor,
+                                       targets: List[torch.nn.Module],
+                                       eigen_smooth: bool = False) -> np.ndarray:
+        transforms = tta.Compose(
+            [
+                tta.HorizontalFlip(),
+                tta.Multiply(factors=[0.9, 1, 1.1]),
+            ]
+        )
+        cams = []
+        for transform in transforms:
+            augmented_tensor = transform.augment_image(input_tensor)
+            cam = self.forward(augmented_tensor,
+                               targets,
+                               eigen_smooth)
+
+            # The ttach library expects a tensor of size BxCxHxW
+            cam = cam[:, None, :, :]
+            cam = torch.from_numpy(cam)
+            cam = transform.deaugment_mask(cam)
+
+            # Back to numpy float32, HxW
+            cam = cam.numpy()
+            cam = cam[:, 0, :, :]
+            cams.append(cam)
+
+        cam = np.mean(np.float32(cams), axis=0)
+        return cam
+
+    def __call__(self,
+                 input_tensor: torch.Tensor,
+                 targets: List[torch.nn.Module] = None,
+                 target_size=None,
+                 aug_smooth: bool = False,
+                 eigen_smooth: bool = False) -> np.ndarray:
+
+        # Smooth the CAM result with test time augmentation
+        if aug_smooth is True:
+            return self.forward_augmentation_smoothing(
+                input_tensor, targets, eigen_smooth)
+
+        return self.forward(input_tensor,
+                            targets, target_size,eigen_smooth)
+
+    def __del__(self):
+        self.activations_and_grads.release()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        self.activations_and_grads.release()
+        if isinstance(exc_value, IndexError):
+            # Handle IndexError here...
+            print(
+                f"An exception occurred in CAM with block: {exc_type}. Message: {exc_value}")
+            return True

models/dsp/CAM/pytorch_grad_cam/eigen_cam.py

@@ -0,0 +1,23 @@
+from .base_cam import BaseCAM
+from .utils.svd_on_activations import get_2d_projection
+
+# https://arxiv.org/abs/2008.00299
+
+
+class EigenCAM(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        super(EigenCAM, self).__init__(model,
+                                       target_layers,
+                                       use_cuda,
+                                       reshape_transform,
+                                       uses_gradients=False)
+
+    def get_cam_image(self,
+                      input_tensor,
+                      target_layer,
+                      target_category,
+                      activations,
+                      grads,
+                      eigen_smooth):
+        return get_2d_projection(activations)

models/dsp/CAM/pytorch_grad_cam/eigen_grad_cam.py

@@ -0,0 +1,21 @@
+from .base_cam import BaseCAM
+from .utils.svd_on_activations import get_2d_projection
+
+# Like Eigen CAM: https://arxiv.org/abs/2008.00299
+# But multiply the activations x gradients
+
+
+class EigenGradCAM(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        super(EigenGradCAM, self).__init__(model, target_layers, use_cuda,
+                                           reshape_transform)
+
+    def get_cam_image(self,
+                      input_tensor,
+                      target_layer,
+                      target_category,
+                      activations,
+                      grads,
+                      eigen_smooth):
+        return get_2d_projection(grads * activations)

models/dsp/CAM/pytorch_grad_cam/fullgrad_cam.py

@@ -0,0 +1,95 @@
+import numpy as np
+import torch
+from .base_cam import BaseCAM
+from .utils.find_layers import find_layer_predicate_recursive
+from .utils.svd_on_activations import get_2d_projection
+from .utils.image import scale_accross_batch_and_channels, scale_cam_image
+
+# https://arxiv.org/abs/1905.00780
+
+
+class FullGrad(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        if len(target_layers) > 0:
+            print(
+                "Warning: target_layers is ignored in FullGrad. All bias layers will be used instead")
+
+        def layer_with_2D_bias(layer):
+            bias_target_layers = [torch.nn.Conv2d, torch.nn.BatchNorm2d]
+            if type(layer) in bias_target_layers and layer.bias is not None:
+                return True
+            return False
+        target_layers = find_layer_predicate_recursive(
+            model, layer_with_2D_bias)
+        super(
+            FullGrad,
+            self).__init__(
+            model,
+            target_layers,
+            use_cuda,
+            reshape_transform,
+            compute_input_gradient=True)
+        self.bias_data = [self.get_bias_data(
+            layer).cpu().numpy() for layer in target_layers]
+
+    def get_bias_data(self, layer):
+        # Borrowed from official paper impl:
+        # https://github.com/idiap/fullgrad-saliency/blob/master/saliency/tensor_extractor.py#L47
+        if isinstance(layer, torch.nn.BatchNorm2d):
+            bias = - (layer.running_mean * layer.weight
+                      / torch.sqrt(layer.running_var + layer.eps)) + layer.bias
+            return bias.data
+        else:
+            return layer.bias.data
+
+    def compute_cam_per_layer(
+            self,
+            input_tensor,
+            target_category,
+            eigen_smooth):
+        input_grad = input_tensor.grad.data.cpu().numpy()
+        grads_list = [g.cpu().data.numpy() for g in
+                      self.activations_and_grads.gradients]
+        cam_per_target_layer = []
+        target_size = self.get_target_width_height(input_tensor)
+
+        gradient_multiplied_input = input_grad * input_tensor.data.cpu().numpy()
+        gradient_multiplied_input = np.abs(gradient_multiplied_input)
+        gradient_multiplied_input = scale_accross_batch_and_channels(
+            gradient_multiplied_input,
+            target_size)
+        cam_per_target_layer.append(gradient_multiplied_input)
+
+        # Loop over the saliency image from every layer
+        assert(len(self.bias_data) == len(grads_list))
+        for bias, grads in zip(self.bias_data, grads_list):
+            bias = bias[None, :, None, None]
+            # In the paper they take the absolute value,
+            # but possibily taking only the positive gradients will work
+            # better.
+            bias_grad = np.abs(bias * grads)
+            result = scale_accross_batch_and_channels(
+                bias_grad, target_size)
+            result = np.sum(result, axis=1)
+            cam_per_target_layer.append(result[:, None, :])
+        cam_per_target_layer = np.concatenate(cam_per_target_layer, axis=1)
+        if eigen_smooth:
+            # Resize to a smaller image, since this method typically has a very large number of channels,
+            # and then consumes a lot of memory
+            cam_per_target_layer = scale_accross_batch_and_channels(
+                cam_per_target_layer, (target_size[0] // 8, target_size[1] // 8))
+            cam_per_target_layer = get_2d_projection(cam_per_target_layer)
+            cam_per_target_layer = cam_per_target_layer[:, None, :, :]
+            cam_per_target_layer = scale_accross_batch_and_channels(
+                cam_per_target_layer,
+                target_size)
+        else:
+            cam_per_target_layer = np.sum(
+                cam_per_target_layer, axis=1)[:, None, :]
+
+        return cam_per_target_layer
+
+    def aggregate_multi_layers(self, cam_per_target_layer):
+        result = np.sum(cam_per_target_layer, axis=1)
+        return scale_cam_image(result)

models/dsp/CAM/pytorch_grad_cam/grad_cam.py

@@ -0,0 +1,24 @@
+import torch
+import numpy as np
+from .base_cam import BaseCAM
+
+
+class GradCAM(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        super(
+            GradCAM,
+            self).__init__(
+            model,
+            target_layers,
+            use_cuda,
+            reshape_transform)
+
+    def get_cam_weights(self,
+                        input_tensor,
+                        target_layer,
+                        target_category,
+                        activations,
+                        grads):
+        # return np.mean(grads, axis=(2, 3))
+        return torch.mean(grads, dim=(2, 3))

models/dsp/CAM/pytorch_grad_cam/grad_cam_plusplus.py

@@ -0,0 +1,32 @@
+import numpy as np
+from .base_cam import BaseCAM
+
+# https://arxiv.org/abs/1710.11063
+
+
+class GradCAMPlusPlus(BaseCAM):
+    def __init__(self, model, target_layers, use_cuda=False,
+                 reshape_transform=None):
+        super(GradCAMPlusPlus, self).__init__(model, target_layers, use_cuda,
+                                              reshape_transform)
+
+    def get_cam_weights(self,
+                        input_tensor,
+                        target_layers,
+                        target_category,
+                        activations,
+                        grads):
+        grads_power_2 = grads**2
+        grads_power_3 = grads_power_2 * grads
+        # Equation 19 in https://arxiv.org/abs/1710.11063
+        sum_activations = np.sum(activations, axis=(2, 3))
+        eps = 0.000001
+        aij = grads_power_2 / (2 * grads_power_2 +
+                               sum_activations[:, :, None, None] * grads_power_3 + eps)
+        # Now bring back the ReLU from eq.7 in the paper,
+        # And zero out aijs where the activations are 0
+        aij = np.where(grads != 0, aij, 0)
+
+        weights = np.maximum(grads, 0) * aij
+        weights = np.sum(weights, axis=(2, 3))
+        return weights

models/dsp/CAM/pytorch_grad_cam/guided_backprop.py

@@ -0,0 +1,100 @@
+import numpy as np
+import torch
+from torch.autograd import Function
+from .utils.find_layers import replace_all_layer_type_recursive
+
+
+class GuidedBackpropReLU(Function):
+    @staticmethod
+    def forward(self, input_img):
+        positive_mask = (input_img > 0).type_as(input_img)
+        output = torch.addcmul(
+            torch.zeros(
+                input_img.size()).type_as(input_img),
+            input_img,
+            positive_mask)
+        self.save_for_backward(input_img, output)
+        return output
+
+    @staticmethod
+    def backward(self, grad_output):
+        input_img, output = self.saved_tensors
+        grad_input = None
+
+        positive_mask_1 = (input_img > 0).type_as(grad_output)
+        positive_mask_2 = (grad_output > 0).type_as(grad_output)
+        grad_input = torch.addcmul(
+            torch.zeros(
+                input_img.size()).type_as(input_img),
+            torch.addcmul(
+                torch.zeros(
+                    input_img.size()).type_as(input_img),
+                grad_output,
+                positive_mask_1),
+            positive_mask_2)
+        return grad_input
+
+
+class GuidedBackpropReLUasModule(torch.nn.Module):
+    def __init__(self):
+        super(GuidedBackpropReLUasModule, self).__init__()
+
+    def forward(self, input_img):
+        return GuidedBackpropReLU.apply(input_img)
+
+
+class GuidedBackpropReLUModel:
+    def __init__(self, model, use_cuda):
+        self.model = model
+        self.model.eval()
+        self.cuda = use_cuda
+        if self.cuda:
+            self.model = self.model.cuda()
+
+    def forward(self, input_img):
+        return self.model(input_img)
+
+    def recursive_replace_relu_with_guidedrelu(self, module_top):
+
+        for idx, module in module_top._modules.items():
+            self.recursive_replace_relu_with_guidedrelu(module)
+            if module.__class__.__name__ == 'ReLU':
+                module_top._modules[idx] = GuidedBackpropReLU.apply
+        print("b")
+
+    def recursive_replace_guidedrelu_with_relu(self, module_top):
+        try:
+            for idx, module in module_top._modules.items():
+                self.recursive_replace_guidedrelu_with_relu(module)
+                if module == GuidedBackpropReLU.apply:
+                    module_top._modules[idx] = torch.nn.ReLU()
+        except BaseException:
+            pass
+
+    def __call__(self, input_img, target_category=None):
+        replace_all_layer_type_recursive(self.model,
+                                         torch.nn.ReLU,
+                                         GuidedBackpropReLUasModule())
+
+        if self.cuda:
+            input_img = input_img.cuda()
+
+        input_img = input_img.requires_grad_(True)
+
+        output = self.forward(input_img)
+
+        if target_category is None:
+            target_category = np.argmax(output.cpu().data.numpy())
+
+        loss = output[0, target_category]
+        loss.backward(retain_graph=True)
+
+        output = input_img.grad.cpu().data.numpy()
+        output = output[0, :, :, :]
+        output = output.transpose((1, 2, 0))
+
+        replace_all_layer_type_recursive(self.model,
+                                         GuidedBackpropReLUasModule,
+                                         torch.nn.ReLU())
+
+        return output

models/dsp/CAM/pytorch_grad_cam/layer_cam.py

@@ -0,0 +1,36 @@
+import numpy as np
+from .base_cam import BaseCAM
+from .utils.svd_on_activations import get_2d_projection
+
+# https://ieeexplore.ieee.org/document/9462463
+
+
+class LayerCAM(BaseCAM):
+    def __init__(
+            self,
+            model,
+            target_layers,
+            use_cuda=False,
+            reshape_transform=None):
+        super(
+            LayerCAM,
+            self).__init__(
+            model,
+            target_layers,
+            use_cuda,
+            reshape_transform)
+
+    def get_cam_image(self,
+                      input_tensor,
+                      target_layer,
+                      target_category,
+                      activations,
+                      grads,
+                      eigen_smooth):
+        spatial_weighted_activations = np.maximum(grads, 0) * activations
+
+        if eigen_smooth:
+            cam = get_2d_projection(spatial_weighted_activations)
+        else:
+            cam = spatial_weighted_activations.sum(axis=1)
+        return cam

models/dsp/CAM/pytorch_grad_cam/score_cam.py

@@ -0,0 +1,63 @@
+import torch
+import tqdm
+from .base_cam import BaseCAM
+
+
+class ScoreCAM(BaseCAM):
+    def __init__(
+            self,
+            model,
+            target_layers,
+            use_cuda=False,
+            reshape_transform=None):
+        super(ScoreCAM, self).__init__(model,
+                                       target_layers,
+                                       use_cuda,
+                                       reshape_transform=reshape_transform,
+                                       uses_gradients=False)
+
+        if len(target_layers) > 0:
+            print("Warning: You are using ScoreCAM with target layers, "
+                  "however ScoreCAM will ignore them.")
+
+    def get_cam_weights(self,
+                        input_tensor,
+                        target_layer,
+                        targets,
+                        activations,
+                        grads):
+        with torch.no_grad():
+            upsample = torch.nn.UpsamplingBilinear2d(
+                size=input_tensor.shape[-2:])
+            activation_tensor = torch.from_numpy(activations)
+            if self.cuda:
+                activation_tensor = activation_tensor.cuda()
+
+            upsampled = upsample(activation_tensor)
+
+            maxs = upsampled.view(upsampled.size(0),
+                                  upsampled.size(1), -1).max(dim=-1)[0]
+            mins = upsampled.view(upsampled.size(0),
+                                  upsampled.size(1), -1).min(dim=-1)[0]
+
+            maxs, mins = maxs[:, :, None, None], mins[:, :, None, None]
+            upsampled = (upsampled - mins) / (maxs - mins)
+
+            input_tensors = input_tensor[:, None,
+                                         :, :] * upsampled[:, :, None, :, :]
+
+            if hasattr(self, "batch_size"):
+                BATCH_SIZE = self.batch_size
+            else:
+                BATCH_SIZE = 16
+
+            scores = []
+            for target, tensor in zip(targets, input_tensors):
+                for i in tqdm.tqdm(range(0, tensor.size(0), BATCH_SIZE)):
+                    batch = tensor[i: i + BATCH_SIZE, :]
+                    outputs = [target(o).cpu().item() for o in self.model(batch)]
+                    scores.extend(outputs)
+            scores = torch.Tensor(scores)
+            scores = scores.view(activations.shape[0], activations.shape[1])
+            weights = torch.nn.Softmax(dim=-1)(scores).numpy()
+            return weights

models/dsp/CAM/pytorch_grad_cam/utils/__init__.py

@@ -0,0 +1,4 @@
+from .image import deprocess_image
+from .svd_on_activations import get_2d_projection
+from . import model_targets
+from . import reshape_transforms

models/dsp/CAM/pytorch_grad_cam/utils/find_layers.py

@@ -0,0 +1,30 @@
+def replace_layer_recursive(model, old_layer, new_layer):
+    for name, layer in model._modules.items():
+        if layer == old_layer:
+            model._modules[name] = new_layer
+            return True
+        elif replace_layer_recursive(layer, old_layer, new_layer):
+            return True
+    return False
+
+
+def replace_all_layer_type_recursive(model, old_layer_type, new_layer):
+    for name, layer in model._modules.items():
+        if isinstance(layer, old_layer_type):
+            model._modules[name] = new_layer
+        replace_all_layer_type_recursive(layer, old_layer_type, new_layer)
+
+
+def find_layer_types_recursive(model, layer_types):
+    def predicate(layer):
+        return type(layer) in layer_types
+    return find_layer_predicate_recursive(model, predicate)
+
+
+def find_layer_predicate_recursive(model, predicate):
+    result = []
+    for name, layer in model._modules.items():
+        if predicate(layer):
+            result.append(layer)
+        result.extend(find_layer_predicate_recursive(layer, predicate))
+    return result

models/dsp/CAM/pytorch_grad_cam/utils/image.py

@@ -0,0 +1,90 @@
+import cv2
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torchvision.transforms import Compose, Normalize, ToTensor
+
+
+def preprocess_image(img: np.ndarray, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) -> torch.Tensor:
+    preprocessing = Compose([
+        ToTensor(),
+        Normalize(mean=mean, std=std)
+    ])
+    return preprocessing(img.copy()).unsqueeze(0)
+
+
+def deprocess_image(img):
+    """ see https://github.com/jacobgil/keras-grad-cam/blob/master/grad-cam.py#L65 """
+    img = img - np.mean(img)
+    img = img / (np.std(img) + 1e-5)
+    img = img * 0.1
+    img = img + 0.5
+    img = np.clip(img, 0, 1)
+    return np.uint8(img * 255)
+
+
+def show_cam_on_image(img: np.ndarray,
+                      mask: np.ndarray,
+                      use_rgb: bool = False,
+                      colormap: int = cv2.COLORMAP_JET) -> np.ndarray:
+    """ This function overlays the cam mask on the image as an heatmap.
+    By default the heatmap is in BGR format.
+
+    :param img: The base image in RGB or BGR format.
+    :param mask: The cam mask.
+    :param use_rgb: Whether to use an RGB or BGR heatmap, this should be set to True if 'img' is in RGB format.
+    :param colormap: The OpenCV colormap to be used.
+    :returns: The default image with the cam overlay.
+    """
+    heatmap = cv2.applyColorMap(np.uint8(255 * mask), colormap)
+    if use_rgb:
+        heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+    heatmap = np.float32(heatmap) / 255
+
+    if np.max(img) > 1:
+        raise Exception(
+            "The input image should np.float32 in the range [0, 1]")
+
+    cam = heatmap + img
+    cam = cam / np.max(cam)
+    return np.uint8(255 * cam)
+
+def scale_cam_image(cam, target_size=None):
+    result = []
+    for img in cam:
+        img = img - np.min(img)
+        img = img / (1e-7 + np.max(img))
+        if target_size is not None:
+            img = cv2.resize(img, target_size)
+        result.append(img)
+    result = np.float32(result)
+
+    return result
+
+def scale_cam_image_torch(cam: torch.Tensor, target_size=None):
+    if cam.ndim == 3:
+        cam = cam.unsqueeze(1)
+
+    # normalize per image
+    cam_min = cam.amin(dim=(2, 3), keepdim=True)
+    cam_max = cam.amax(dim=(2, 3), keepdim=True)
+    cam = cam - cam_min
+    cam = cam / (cam_max + 1e-7)
+
+    # resize if needed
+    if target_size is not None:
+        cam = F.interpolate(cam, size=target_size, mode="bilinear", align_corners=False)
+
+    return cam.squeeze(1)
+
+def scale_accross_batch_and_channels(tensor, target_size):
+    batch_size, channel_size = tensor.shape[:2]
+    reshaped_tensor = tensor.reshape(
+        batch_size * channel_size, *tensor.shape[2:])
+    result = scale_cam_image(reshaped_tensor, target_size)
+    result = result.reshape(
+        batch_size,
+        channel_size,
+        target_size[1],
+        target_size[0])
+    return result