diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..62d3c63488fb2cc487ecb049b6cc8fc638757776 --- /dev/null +++ b/.gitignore @@ -0,0 +1,26 @@ +__pycache__/ +*.py[cod] +*.egg-info/ +.eggs/ +dist/ +build/ +.DS_Store +logs/ +*.log +coco_log/ +.ipynb_checkpoints/ +_figs/ +_paper_extract.txt + +# paper PDF — do not upload +FlexICM.pdf +*.pdf + +# real weights (keep PLACEHOLDER text files and README) +checkpoints/**/*.pth.tar +checkpoints/**/*.tar +checkpoints/**/*.pkl +checkpoints/**/*.pth +!checkpoints/**/PLACEHOLDER* +!checkpoints/**/*.txt +!checkpoints/README.md diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..d590c0f725117d5037b64e7b182255a7fa6824c4 --- /dev/null +++ b/README.md @@ -0,0 +1,316 @@ +# FlexICM: A Flexible Image Coding for Machines Framework + +Official codebase for the paper **FlexICM: A Flexible Image Coding for Machines Framework** (Tianma Shen, Ying Liu). + +Built on the **TIC (Transformer-based Image Compression)** base codec, this repository implements: + +- **TAIC (Base Layer)**: five single-task codecs that decode task intermediate features `h` **without** full image reconstruction +- **C-TAIC (Extension Layer)**: three multi-task scenarios that condition on the base-layer latent \(\hat{y}_b\) via cross-attention + +## Five Tasks and Three Scenarios + +### TAIC (five task codecs) + +| Task | Teacher / Task Network | Feature Alignment | Metric | +|------|------------------------|-------------------|--------| +| Object Detection | Faster R-CNN + **Swin-B** | FPN `P2..P6` (Eq. 2) | mAP-bbox | +| Semantic Segmentation | UPerNet + **Swin-B** | FPN `P2..P6` | mIoU | +| Instance Segmentation | Mask R-CNN + **Swin-B** | FPN `P2..P6` | mAP-mask | +| Panoptic Segmentation | MaskFormer + **Swin-B** | Stages `F1..F4` (Eq. 3) | PQ | +| Pose Estimation | **HigherHRNet** | Stages `F1..F4` | mAP-OKS | + +### C-TAIC (three scenarios) + +| Scenario | Base Layer | Extension Layer | +|----------|------------|-----------------| +| **s1** | Object Detection | Instance Segmentation | +| **s2** | Semantic Segmentation | Panoptic Segmentation | +| **s3** | Object Detection | Pose Estimation | + +--- + +## Environment Setup + +> **Important:** Codec training **requires** task networks (teachers) to be available. +> The loss \(D\) is computed from frozen teacher features, so you cannot train TAIC / C-TAIC +> with only the codec packages. Install the teacher stack in **Task networks (teachers)** before the first training run. + +### Recommended environment + +- Ubuntu / RHEL, **CUDA 11.7+**, single **NVIDIA A100** (paper setting) +- Python **3.8–3.10** +- PyTorch **≥ 1.12** (2.0+ recommended) + +```bash +conda create -n flexicm python=3.9 -y +conda activate flexicm +pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118 +pip install -r requirements.txt +``` + +### Core codec dependencies + +| Package | Role | +|---------|------| +| `compressai` | EntropyBottleneck / GaussianConditional / conv-deconv | +| `timm` | **Required** Swin-B teacher backbone for feature alignment | +| `PyYAML` | Training configs | + +### Task networks (teachers) — **required before training** + +Teachers are already implemented in `flexicm/tasks/` and are constructed automatically by +`scripts/train_taic.py` / `scripts/train_ctaic.py` via `build_teacher(...)`. +You still must install their runtime dependencies and allow pretrained weights to download. + +| Task | Teacher used in training | What you need installed | +|------|--------------------------|-------------------------| +| Detection / Instance / Semantic / Panoptic | Swin-B backbone (+ FPN or stages) via `timm` | `timm` (from `requirements.txt`); first run downloads ImageNet-pretrained Swin-B | +| Pose | HigherHRNet-style HRNet stem (original HRNet, not Swin) | Implemented in-repo; no extra package beyond PyTorch | + +Checklist before training: + +1. `pip install -r requirements.txt` (includes `timm`) +2. Machine can reach the internet **or** you have cached `timm` Swin-B weights (for the four Swin tasks) +3. Verify teachers import cleanly: + +```bash +python -c "from flexicm.tasks import build_teacher; build_teacher('detection'); print('teachers ok')" +``` + +Without a working teacher, training will fail when computing the feature-alignment term \(D\). + +### Task heads for metric evaluation + +To evaluate paper metrics (mAP / mIoU / PQ / OKS) with full task heads, also install: + +```bash +pip install -U openmim +mim install mmengine mmcv +mim install mmdet mmsegmentation mmpose +# or Detectron2 (alternative for detection / instance evaluation) +``` + +Recommended official weights (same model families as the paper): + +- **Faster / Mask R-CNN + Swin-B**: MMDetection Model Zoo +- **UPerNet + Swin-B**: MMSegmentation Model Zoo +- **MaskFormer + Swin-B**: MMDetection / Mask2Former +- **HigherHRNet**: MMPose Model Zoo (**HRNet backbone**) + +These full heads are **not** required to start codec training; they are for final rate–accuracy evaluation. + +--- + +## Repository Layout + +``` +FlexICM/ +├── FlexICM.pdf # paper +├── requirements.txt +├── README.md +├── configs/ +│ ├── taic/ # five single-task configs +│ │ ├── detection.yaml +│ │ ├── semantic.yaml +│ │ ├── instance.yaml +│ │ ├── panoptic.yaml +│ │ └── pose.yaml +│ └── ctaic/ # three multi-task scenarios +│ ├── s1_det_instance.yaml +│ ├── s2_sem_panoptic.yaml +│ └── s3_det_pose.yaml +├── scripts/ +│ ├── download_base_codecs.sh +│ ├── train_taic.py +│ └── train_ctaic.py +├── flexicm/ +│ ├── models/ # TAIC / C-TAIC / SFMA / TaskConnector / Conditional +│ ├── layers/ # RSTB / WindowAttention (same lineage as AdaptiveICMH) +│ ├── tasks/ # teachers + feature-alignment losses +│ ├── data/ # COCO / COCO-WholeBody image loading +│ └── utils/ +└── checkpoints/ # placeholder tree for base / TAIC / C-TAIC weights + # see checkpoints/README.md +``` + +Eval configs (stub until full metrics are implemented): `configs/eval/`. +Eval entry points: `scripts/eval_taic.py`, `scripts/eval_ctaic.py` (currently only check that real checkpoints replaced `PLACEHOLDER` files). + +--- + +## Dataset Preparation + +### COCO-2017 (detection / instance / semantic / panoptic) + +```text +/data/coco2017/ +├── train2017/ +├── val2017/ +└── annotations/ + ├── instances_train2017.json + ├── instances_val2017.json + ├── panoptic_train2017.json + ├── panoptic_val2017.json + ├── panoptic_train2017/ # PNG + ├── panoptic_val2017/ + ├── stuff_train2017.json # semantic / stuff (if used) + └── stuff_val2017.json +``` + +Download: + +```bash +# images +wget http://images.cocodataset.org/zips/train2017.zip +wget http://images.cocodataset.org/zips/val2017.zip +# annotations +wget http://images.cocodataset.org/annotations/annotations_trainval2017.zip +wget http://images.cocodataset.org/annotations/panoptic_annotations_trainval2017.zip +``` + +Set in the corresponding YAML: + +```yaml +dataset_path: "/data/coco2017" +``` + +### COCO-WholeBody (pose estimation) + +Pose uses the same COCO `train2017/val2017` images plus WholeBody keypoint annotations: + +- Project page: [COCO-WholeBody](https://github.com/jin-s13/COCO-WholeBody) +- Place JSON files under `annotations/`; evaluate with MMPose HigherHRNet + WholeBody configs + +Codec **training** only needs images for feature alignment, so `train2017` images are sufficient for that stage. + +### Data processing (aligned with task networks) + +The paper requires **codec training preprocessing to match task-network preprocessing**. Defaults in this repo: + +1. **Codec input**: RGB, `ToTensor()` → `[0,1]`; training uses `Resize → RandomCrop(256) → RandomHorizontalFlip` +2. **Inside the teacher**: ImageNet mean/std normalization (consistent with Swin / HRNet pretraining) +3. **Spatial alignment**: TIC requires spatial size divisible by **256** (256 crop for training; pad at inference) + +If you use official MMDet/MMSeg pipelines (short-side resize, normalization, etc.), ensure: + +- Teacher feature extraction uses the **same normalize / resize logic** as that task network +- Codec and teacher see geometrically consistent tensors (same crop / same pad) + +Edit points: `flexicm/data/datasets.py`, `flexicm/tasks/swin_teacher.py`, `flexicm/tasks/__init__.py` (HigherHRNet). + +--- + +## Base Codec (TIC) Checkpoints + +The paper uses the same TIC pretrained weights as AdaptiveICMH / TransTIC: + +| Quality | λ (paper) | Checkpoint | +|:-------:|:---------:|------------| +| 1 | 0.0035 | [base_codec_1](https://github.com/NYCU-MAPL/TransTIC/releases/download/v1.0/base_codec_1.pth.tar) | +| 2 | 0.0067 | [base_codec_2](https://github.com/NYCU-MAPL/TransTIC/releases/download/v1.0/base_codec_2.pth.tar) | +| 3 | 0.0130 | [base_codec_3](https://github.com/NYCU-MAPL/TransTIC/releases/download/v1.0/base_codec_3.pth.tar) | +| 4 | 0.0250 | [base_codec_4](https://github.com/NYCU-MAPL/TransTIC/releases/download/v1.0/base_codec_4.pth.tar) | + +```bash +bash scripts/download_base_codecs.sh +# downloads into checkpoints/base_codec/base_codec_{1,2,3,4}.pth.tar +``` + +Config example: + +```yaml +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +quality_level: 1 +lmbda: 0.0035 +``` + +For each bitrate point, switch the matching `base_codec_k` and `lmbda`. + +Trained TAIC / C-TAIC weights for eval should be placed under `checkpoints/taic/` and +`checkpoints/ctaic/` (see `checkpoints/README.md`). Until then, each quality folder +contains a `PLACEHOLDER` file. + + +--- + +## Training + +Paper settings: + +- Optimizer: **AdamW**, `lr=1e-4` +- TAIC: `batch_size=80`, `epochs=35` +- C-TAIC: `batch_size=40`, `epochs=40` +- `λ ∈ {0.0035, 0.0067, 0.0130, 0.0250}` + +> If GPU memory is insufficient, reduce `batch_size` (optionally use gradient accumulation to approximate the paper effective batch). + +### Train five TAIC models + +```bash +# edit dataset_path / base_codec / lmbda / gpu_id in configs/taic/*.yaml as needed +python scripts/train_taic.py -c configs/taic/detection.yaml +python scripts/train_taic.py -c configs/taic/semantic.yaml +python scripts/train_taic.py -c configs/taic/instance.yaml +python scripts/train_taic.py -c configs/taic/panoptic.yaml +python scripts/train_taic.py -c configs/taic/pose.yaml +``` + +Trainable modules: **encoder SFMA + Task Connector**; TIC trunk is frozen. + +### Train three C-TAIC scenarios + +Requires a trained **base TAIC** checkpoint (to provide \(\hat{y}_b\)) and Stage-1 weights for the extension task. + +```bash +# ---- s1: det → instance ---- +python scripts/train_ctaic.py -c configs/ctaic/s1_det_instance.yaml --stage 1 +python scripts/train_ctaic.py -c configs/ctaic/s1_det_instance.yaml --stage 2 + +# ---- s2: semantic → panoptic ---- +python scripts/train_ctaic.py -c configs/ctaic/s2_sem_panoptic.yaml --stage 1 +python scripts/train_ctaic.py -c configs/ctaic/s2_sem_panoptic.yaml --stage 2 + +# ---- s3: det → pose ---- +python scripts/train_ctaic.py -c configs/ctaic/s3_det_pose.yaml --stage 1 +python scripts/train_ctaic.py -c configs/ctaic/s3_det_pose.yaml --stage 2 +``` + +Stage meanings: + +| Stage | Mode | Trainable modules | `ŷ_b` | +|:-----:|------|-------------------|-------| +| 1 | TAIC mode | SFMA + Task Connector | not used | +| 2 | C-TAIC mode | Prompt Generator + Condition Generator | from frozen base TAIC AD output | + +Check these config fields: + +```yaml +base_taic_checkpoint: # trained TAIC for the base task +taic_init: # extension-task TAIC (optional Stage-1 init) +stage1_checkpoint: # Stage-1 result loaded in Stage 2 +``` + +--- + +## Code Map to the Paper + +| Paper component | Code location | +|-----------------|---------------| +| SFMA | `flexicm/models/sfma.py` | +| Task Connector | `flexicm/models/task_connector.py` | +| TAIC | `flexicm/models/taic.py` | +| C-TAIC + two-stage freeze | `flexicm/models/ctaic.py` | +| Prompt / Mask / Cd | `flexicm/models/conditional.py` | +| Cross-attention (Q from features; K/V include prompts) | `flexicm/models/cross_attention.py` | +| \(R+\lambda D\) | `flexicm/tasks/losses.py` | +| Five teachers | `flexicm/tasks/__init__.py` | + + +--- + +## Citation + +If you use this code or the paper, please cite FlexICM and acknowledge the base works: + +- TIC: Lu et al., Transformer-based Image Compression +- TransTIC / AdaptiveICMH: task-adaptive SFMA tuning diff --git a/checkpoints/README.md b/checkpoints/README.md new file mode 100644 index 0000000000000000000000000000000000000000..cd0db62837d0848c3ee5fb3493b1d418b480e34e --- /dev/null +++ b/checkpoints/README.md @@ -0,0 +1,64 @@ +# Checkpoints Layout (Placeholders) + +Put pretrained / trained weights here before running **test / eval**. +Training still writes to `logs/` by default; after training, copy (or symlink) best +checkpoints into this tree so eval configs have a stable path. + +> Files named `PLACEHOLDER` are not real weights. Replace each with the matching +> `.pth.tar` checkpoint, then update or keep the path expected by eval scripts. + +## Directory map + +```text +checkpoints/ +├── base_codec/ # frozen TIC (TransTIC / AdaptiveICMH) +│ ├── base_codec_1.pth.tar # λ = 0.0035 +│ ├── base_codec_2.pth.tar # λ = 0.0067 +│ ├── base_codec_3.pth.tar # λ = 0.0130 +│ └── base_codec_4.pth.tar # λ = 0.0250 +│ +├── taic/ # five single-task TAIC codecs +│ ├── detection/{1,2,3,4}/checkpoint_best_loss.pth.tar +│ ├── semantic/{1,2,3,4}/checkpoint_best_loss.pth.tar +│ ├── instance/{1,2,3,4}/checkpoint_best_loss.pth.tar +│ ├── panoptic/{1,2,3,4}/checkpoint_best_loss.pth.tar +│ └── pose/{1,2,3,4}/checkpoint_best_loss.pth.tar +│ +└── ctaic/ # three multi-task scenarios + ├── s1_det_instance/ + │ ├── stage1/{1,2,3,4}/checkpoint_best_loss.pth.tar + │ └── stage2/{1,2,3,4}/checkpoint_best_loss.pth.tar + ├── s2_sem_panoptic/ + │ ├── stage1/{1,2,3,4}/checkpoint_best_loss.pth.tar + │ └── stage2/{1,2,3,4}/checkpoint_best_loss.pth.tar + └── s3_det_pose/ + ├── stage1/{1,2,3,4}/checkpoint_best_loss.pth.tar + └── stage2/{1,2,3,4}/checkpoint_best_loss.pth.tar +``` + +Quality folders `{1,2,3,4}` match paper λ / TIC quality levels. + +## Download base TIC codecs + +```bash +bash scripts/download_base_codecs.sh +# downloads into checkpoints/base_codec/ +``` + +## After training: copy into placeholders + +```bash +# example: TAIC detection, quality 1 +cp logs/taic_detection/1/checkpoint_best_loss.pth.tar \ + checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar + +# example: C-TAIC s1 stage2, quality 1 +cp logs/ctaic_s1_stage2/1/checkpoint_best_loss.pth.tar \ + checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_best_loss.pth.tar +``` + +## Eval configs + +See `configs/eval/` — they point to these placeholder paths. +Eval scripts will refuse to run if a `PLACEHOLDER` file is still present +or if the `.pth.tar` is missing. diff --git a/checkpoints/base_codec/NOTE.txt b/checkpoints/base_codec/NOTE.txt new file mode 100644 index 0000000000000000000000000000000000000000..3e35f8267fe1e54844bbb481277780a142b33648 --- /dev/null +++ b/checkpoints/base_codec/NOTE.txt @@ -0,0 +1,2 @@ +Legacy config paths used ./checkpoints/base_codec_k.pth.tar. +Prefer checkpoints/base_codec/base_codec_k.pth.tar after download. diff --git a/checkpoints/base_codec/PLACEHOLDER_base_codec_1.pth.tar.txt b/checkpoints/base_codec/PLACEHOLDER_base_codec_1.pth.tar.txt new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/base_codec/PLACEHOLDER_base_codec_1.pth.tar.txt @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/base_codec/PLACEHOLDER_base_codec_2.pth.tar.txt b/checkpoints/base_codec/PLACEHOLDER_base_codec_2.pth.tar.txt new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/base_codec/PLACEHOLDER_base_codec_2.pth.tar.txt @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/base_codec/PLACEHOLDER_base_codec_3.pth.tar.txt b/checkpoints/base_codec/PLACEHOLDER_base_codec_3.pth.tar.txt new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/base_codec/PLACEHOLDER_base_codec_3.pth.tar.txt @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/base_codec/PLACEHOLDER_base_codec_4.pth.tar.txt b/checkpoints/base_codec/PLACEHOLDER_base_codec_4.pth.tar.txt new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/base_codec/PLACEHOLDER_base_codec_4.pth.tar.txt @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage1/1/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage1/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage1/2/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage1/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage1/3/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage1/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage1/4/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage1/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage2/1/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage2/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage2/2/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage2/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage2/3/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage2/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s1_det_instance/stage2/4/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s1_det_instance/stage2/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/1/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage1/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/2/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage1/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/3/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage1/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/4/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage1/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/1/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage2/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/2/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage2/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/3/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage2/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/4/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s2_sem_panoptic/stage2/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage1/1/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage1/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage1/2/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage1/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage1/3/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage1/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage1/4/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage1/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage2/1/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage2/2/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage2/3/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/ctaic/s3_det_pose/stage2/4/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/detection/1/PLACEHOLDER b/checkpoints/taic/detection/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/detection/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/detection/2/PLACEHOLDER b/checkpoints/taic/detection/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/detection/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/detection/3/PLACEHOLDER b/checkpoints/taic/detection/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/detection/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/detection/4/PLACEHOLDER b/checkpoints/taic/detection/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/detection/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/instance/1/PLACEHOLDER b/checkpoints/taic/instance/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/instance/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/instance/2/PLACEHOLDER b/checkpoints/taic/instance/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/instance/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/instance/3/PLACEHOLDER b/checkpoints/taic/instance/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/instance/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/instance/4/PLACEHOLDER b/checkpoints/taic/instance/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/instance/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/panoptic/1/PLACEHOLDER b/checkpoints/taic/panoptic/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/panoptic/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/panoptic/2/PLACEHOLDER b/checkpoints/taic/panoptic/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/panoptic/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/panoptic/3/PLACEHOLDER b/checkpoints/taic/panoptic/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/panoptic/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/panoptic/4/PLACEHOLDER b/checkpoints/taic/panoptic/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/panoptic/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/pose/1/PLACEHOLDER b/checkpoints/taic/pose/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/pose/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/pose/2/PLACEHOLDER b/checkpoints/taic/pose/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/pose/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/pose/3/PLACEHOLDER b/checkpoints/taic/pose/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/pose/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/pose/4/PLACEHOLDER b/checkpoints/taic/pose/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/pose/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/semantic/1/PLACEHOLDER b/checkpoints/taic/semantic/1/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/semantic/1/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/semantic/2/PLACEHOLDER b/checkpoints/taic/semantic/2/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/semantic/2/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/semantic/3/PLACEHOLDER b/checkpoints/taic/semantic/3/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/semantic/3/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/checkpoints/taic/semantic/4/PLACEHOLDER b/checkpoints/taic/semantic/4/PLACEHOLDER new file mode 100644 index 0000000000000000000000000000000000000000..c8768731085e17b2e27011b46e84e4423af2019f --- /dev/null +++ b/checkpoints/taic/semantic/4/PLACEHOLDER @@ -0,0 +1,2 @@ +PLACEHOLDER: replace this file with the real checkpoint (.pth.tar). +See checkpoints/README.md for the expected filename and training copy commands. diff --git a/configs/ctaic/s1_det_instance.yaml b/configs/ctaic/s1_det_instance.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5318dcee77fb74388740e10288eed35a21d8376a --- /dev/null +++ b/configs/ctaic/s1_det_instance.yaml @@ -0,0 +1,24 @@ +# Scenario s1: Object Detection (base) + Instance Segmentation (extension) +root: "logs" +exp_name: "ctaic_s1" +scenario: "s1" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +# trained base-layer TAIC for detection (provides y_b_hat) +base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar" +# optional TAIC init for extension task +taic_init: "./checkpoints/taic/instance/1/checkpoint_best_loss.pth.tar" +stage1_checkpoint: "./checkpoints/ctaic/s1_det_instance/stage1/1/checkpoint_best_loss.pth.tar" +epochs: 40 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 40 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true diff --git a/configs/ctaic/s2_sem_panoptic.yaml b/configs/ctaic/s2_sem_panoptic.yaml new file mode 100644 index 0000000000000000000000000000000000000000..14b6473837c83216a1fddb2499aee9f3d213596b --- /dev/null +++ b/configs/ctaic/s2_sem_panoptic.yaml @@ -0,0 +1,22 @@ +# Scenario s2: Semantic Segmentation (base) + Panoptic Segmentation (extension) +root: "logs" +exp_name: "ctaic_s2" +scenario: "s2" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +base_taic_checkpoint: "./checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar" +taic_init: "./checkpoints/taic/panoptic/1/checkpoint_best_loss.pth.tar" +stage1_checkpoint: "./checkpoints/ctaic/s2_sem_panoptic/stage1/1/checkpoint_best_loss.pth.tar" +epochs: 40 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 40 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true diff --git a/configs/ctaic/s3_det_pose.yaml b/configs/ctaic/s3_det_pose.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d9eaf6ea2f72844fea3a95120c8532f6904ab00e --- /dev/null +++ b/configs/ctaic/s3_det_pose.yaml @@ -0,0 +1,24 @@ +# Scenario s3: Object Detection (base) + Pose Estimation (extension) +root: "logs" +exp_name: "ctaic_s3" +scenario: "s3" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar" +taic_init: "./checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar" +stage1_checkpoint: "./checkpoints/ctaic/s3_det_pose/stage1/1/checkpoint_best_loss.pth.tar" +epochs: 40 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 40 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true +out_channels: 32 +align_mode: "stages" diff --git a/configs/eval/ctaic_s1.yaml b/configs/eval/ctaic_s1.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ede4753ce1f732f7da1acff1544b9b6a9b914bba --- /dev/null +++ b/configs/eval/ctaic_s1.yaml @@ -0,0 +1,12 @@ +# Eval config (stub) — s1: detection (base) + instance (extension) +scenario: "s1" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_best_loss.pth.tar" +base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/ctaic_s2.yaml b/configs/eval/ctaic_s2.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c714fb16d6902b05b8e5d41ea932df5e9f43fdbd --- /dev/null +++ b/configs/eval/ctaic_s2.yaml @@ -0,0 +1,11 @@ +scenario: "s2" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/ctaic/s2_sem_panoptic/stage2/1/checkpoint_best_loss.pth.tar" +base_taic_checkpoint: "./checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/ctaic_s3.yaml b/configs/eval/ctaic_s3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..1545ad20d69c284fa048980f524d3a9398265117 --- /dev/null +++ b/configs/eval/ctaic_s3.yaml @@ -0,0 +1,11 @@ +scenario: "s3" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_best_loss.pth.tar" +base_taic_checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/taic_detection.yaml b/configs/eval/taic_detection.yaml new file mode 100644 index 0000000000000000000000000000000000000000..286bbd839f4e75efeb03a7aa6c2565dbc5812585 --- /dev/null +++ b/configs/eval/taic_detection.yaml @@ -0,0 +1,11 @@ +# Eval config (stub) — replace PLACEHOLDER weights under checkpoints/taic/detection/ +task: "detection" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/taic_instance.yaml b/configs/eval/taic_instance.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0ce3bdddb5f5d5d5537a4f3b0db23d976a9f1534 --- /dev/null +++ b/configs/eval/taic_instance.yaml @@ -0,0 +1,11 @@ +# Eval config (stub) — replace PLACEHOLDER weights under checkpoints/taic/instance/ +task: "instance" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/taic/instance/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/taic_panoptic.yaml b/configs/eval/taic_panoptic.yaml new file mode 100644 index 0000000000000000000000000000000000000000..5854e38fb1be4fca784c60a3d155a436deb3546d --- /dev/null +++ b/configs/eval/taic_panoptic.yaml @@ -0,0 +1,11 @@ +# Eval config (stub) — replace PLACEHOLDER weights under checkpoints/taic/panoptic/ +task: "panoptic" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/taic/panoptic/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/taic_pose.yaml b/configs/eval/taic_pose.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e6a114ff0054a4b9eef4c0adeac75797390d4b5d --- /dev/null +++ b/configs/eval/taic_pose.yaml @@ -0,0 +1,11 @@ +# Eval config (stub) — replace PLACEHOLDER weights under checkpoints/taic/pose/ +task: "pose" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/eval/taic_semantic.yaml b/configs/eval/taic_semantic.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a48ddc41965277e848a08cc1facca7846fd240c8 --- /dev/null +++ b/configs/eval/taic_semantic.yaml @@ -0,0 +1,11 @@ +# Eval config (stub) — replace PLACEHOLDER weights under checkpoints/taic/semantic/ +task: "semantic" +dataset_path: "/data/coco2017" +quality_level: 1 +lmbda: 0.0035 +checkpoint: "./checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +gpu_id: 0 +cuda: true +test_batch_size: 1 +num_workers: 4 diff --git a/configs/taic/detection.yaml b/configs/taic/detection.yaml new file mode 100644 index 0000000000000000000000000000000000000000..69d65d0e634d56d071a2b3caf30b14804f5f1832 --- /dev/null +++ b/configs/taic/detection.yaml @@ -0,0 +1,23 @@ +# FlexICM TAIC - Object Detection (Faster R-CNN + Swin-B) +root: "logs" +exp_name: "taic_detection" +task: "detection" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" # quality matches lmbda mapping +checkpoint: null +epochs: 35 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +# paper lambda set: {0.0035, 0.0067, 0.0130, 0.0250} +lmbda: 0.0035 +num_workers: 8 +batch_size: 80 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true +out_channels: 128 +align_mode: "fpn" diff --git a/configs/taic/instance.yaml b/configs/taic/instance.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0a1d82536b22b01eb7660671583ac4582bc8fc7f --- /dev/null +++ b/configs/taic/instance.yaml @@ -0,0 +1,21 @@ +root: "logs" +exp_name: "taic_instance" +task: "instance" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +checkpoint: null +epochs: 35 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 80 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true +out_channels: 128 +align_mode: "fpn" diff --git a/configs/taic/panoptic.yaml b/configs/taic/panoptic.yaml new file mode 100644 index 0000000000000000000000000000000000000000..e5c34daf66463868701643cde4e379d42a2bc7e9 --- /dev/null +++ b/configs/taic/panoptic.yaml @@ -0,0 +1,21 @@ +root: "logs" +exp_name: "taic_panoptic" +task: "panoptic" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +checkpoint: null +epochs: 35 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 80 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true +out_channels: 128 +align_mode: "stages" diff --git a/configs/taic/pose.yaml b/configs/taic/pose.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9d030eaa93e47c87bb8a5afe023832a20b71af63 --- /dev/null +++ b/configs/taic/pose.yaml @@ -0,0 +1,21 @@ +root: "logs" +exp_name: "taic_pose" +task: "pose" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +checkpoint: null +epochs: 35 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 80 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true +out_channels: 32 +align_mode: "stages" diff --git a/configs/taic/semantic.yaml b/configs/taic/semantic.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a1acba50d8860b21f8efb83b12c061518864e4cc --- /dev/null +++ b/configs/taic/semantic.yaml @@ -0,0 +1,21 @@ +root: "logs" +exp_name: "taic_semantic" +task: "semantic" +dataset_path: "/data/coco2017" +base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" +checkpoint: null +epochs: 35 +learning_rate: 1.0e-4 +gpu_id: 0 +quality_level: 1 +lmbda: 0.0035 +num_workers: 8 +batch_size: 80 +test_batch_size: 1 +patch_size: 256 +cuda: true +save: true +seed: 42 +pretrained_backbone: true +out_channels: 128 +align_mode: "fpn" diff --git a/flexicm/__init__.py b/flexicm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3d528b8eb454360fb0b61d8ca367c4fe477d8154 --- /dev/null +++ b/flexicm/__init__.py @@ -0,0 +1,3 @@ +"""Package init for FlexICM.""" + +__version__ = "1.0.0" diff --git a/flexicm/data/__init__.py b/flexicm/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..aaaef4d3f15faa5333a265d05c7c7dc0dc4d65f5 --- /dev/null +++ b/flexicm/data/__init__.py @@ -0,0 +1,17 @@ +from .datasets import ( + COCOImageDataset, + COCOWholeBodyImageDataset, + ImageFolderDataset, + build_test_transform, + build_train_transform, + collate_keep, +) + +__all__ = [ + "COCOImageDataset", + "COCOWholeBodyImageDataset", + "ImageFolderDataset", + "build_test_transform", + "build_train_transform", + "collate_keep", +] diff --git a/flexicm/data/datasets.py b/flexicm/data/datasets.py new file mode 100644 index 0000000000000000000000000000000000000000..7fcb96c3db9ecd4972d715964de542beaa43a237 --- /dev/null +++ b/flexicm/data/datasets.py @@ -0,0 +1,85 @@ +"""Datasets and preprocessing aligned with task-network training.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Callable, List, Optional, Tuple + +import torch +from PIL import Image +from torch.utils.data import Dataset +from torchvision import transforms + + +IMAGENET_MEAN = [0.485, 0.456, 0.406] +IMAGENET_STD = [0.229, 0.224, 0.225] + + +def build_train_transform(patch_size: int = 256) -> Callable: + """Codec training crops; keep RGB float in [0, 1] (normalization inside teacher).""" + return transforms.Compose( + [ + transforms.Resize(patch_size), + transforms.RandomCrop(patch_size), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), + ] + ) + + +def build_test_transform() -> Callable: + return transforms.ToTensor() + + +class ImageFolderDataset(Dataset): + """Generic image folder (recursive) for codec training.""" + + IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"} + + def __init__(self, root: str, transform: Optional[Callable] = None, list_file: Optional[str] = None): + self.root = root + self.transform = transform + if list_file and os.path.isfile(list_file): + with open(list_file) as f: + rels = [ln.strip() for ln in f if ln.strip()] + self.files = [os.path.join(root, r) if not os.path.isabs(r) else r for r in rels] + else: + self.files = [] + for dirpath, _, filenames in os.walk(root): + for fn in filenames: + if Path(fn).suffix.lower() in self.IMG_EXTS: + self.files.append(os.path.join(dirpath, fn)) + self.files.sort() + if not self.files: + raise FileNotFoundError(f"No images found under {root}") + + def __len__(self): + return len(self.files) + + def __getitem__(self, index): + path = self.files[index] + img = Image.open(path).convert("RGB") + if self.transform: + img = self.transform(img) + return img + + +class COCOImageDataset(ImageFolderDataset): + """COCO images for detection / instance / semantic / panoptic codec training.""" + + def __init__(self, coco_root: str, split: str = "train2017", transform=None, list_file=None): + image_dir = os.path.join(coco_root, split) + super().__init__(image_dir, transform=transform, list_file=list_file) + + +class COCOWholeBodyImageDataset(ImageFolderDataset): + """COCO-WholeBody uses the same COCO images; annotations differ at eval time.""" + + def __init__(self, coco_root: str, split: str = "train2017", transform=None, list_file=None): + image_dir = os.path.join(coco_root, split) + super().__init__(image_dir, transform=transform, list_file=list_file) + + +def collate_keep(batch): + return torch.stack(batch, dim=0) diff --git a/flexicm/layers/__init__.py b/flexicm/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5f078707434ec3670e8a1551aac4ea1c42c16df4 --- /dev/null +++ b/flexicm/layers/__init__.py @@ -0,0 +1,2 @@ +from .layers import * +from .gdn import GDN diff --git a/flexicm/layers/gdn.py b/flexicm/layers/gdn.py new file mode 100644 index 0000000000000000000000000000000000000000..099b987d561034e4ec35184c44e81485f0ccd9c4 --- /dev/null +++ b/flexicm/layers/gdn.py @@ -0,0 +1,121 @@ +# Copyright (c) 2021-2022, InterDigital Communications, Inc +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted (subject to the limitations in the disclaimer +# below) provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# * Neither the name of InterDigital Communications, Inc nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. + +# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY +# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT +# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A +# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR +# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF +# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from torch import Tensor + +from compressai.ops.parametrizers import NonNegativeParametrizer + +__all__ = ["GDN", "GDN1"] + + +class GDN(nn.Module): + r"""Generalized Divisive Normalization layer. + + Introduced in `"Density Modeling of Images Using a Generalized Normalization + Transformation" `_, + by Balle Johannes, Valero Laparra, and Eero P. Simoncelli, (2016). + + .. math:: + + y[i] = \frac{x[i]}{\sqrt{\beta[i] + \sum_j(\gamma[j, i] * x[j]^2)}} + + """ + + def __init__( + self, + in_channels: int, + inverse: bool = False, + beta_min: float = 1e-6, + gamma_init: float = 0.1, + ): + super().__init__() + + beta_min = float(beta_min) + gamma_init = float(gamma_init) + self.inverse = bool(inverse) + + self.beta_reparam = NonNegativeParametrizer(minimum=beta_min) + beta = torch.ones(in_channels) + beta = self.beta_reparam.init(beta) + self.beta = nn.Parameter(beta) + + self.gamma_reparam = NonNegativeParametrizer() + gamma = gamma_init * torch.eye(in_channels) + gamma = self.gamma_reparam.init(gamma) + self.gamma = nn.Parameter(gamma) + + def forward(self, x: Tensor) -> Tensor: + _, C, _, _ = x.size() + + beta = self.beta_reparam(self.beta) + gamma = self.gamma_reparam(self.gamma) + gamma = gamma.reshape(C, C, 1, 1) + norm = F.conv2d(x**2, gamma, beta) + + if self.inverse: + norm = torch.sqrt(norm) + else: + norm = torch.rsqrt(norm) + + out = x * norm + + return out + + +class GDN1(GDN): + r"""Simplified GDN layer. + + Introduced in `"Computationally Efficient Neural Image Compression" + `_, by Johnston Nick, Elad Eban, Ariel + Gordon, and Johannes Ballé, (2019). + + .. math:: + + y[i] = \frac{x[i]}{\beta[i] + \sum_j(\gamma[j, i] * |x[j]|} + + """ + + def forward(self, x: Tensor) -> Tensor: + _, C, _, _ = x.size() + + beta = self.beta_reparam(self.beta) + gamma = self.gamma_reparam(self.gamma) + gamma = gamma.reshape(C, C, 1, 1) + norm = F.conv2d(torch.abs(x), gamma, beta) + + if not self.inverse: + norm = 1.0 / norm + + out = x * norm + + return out diff --git a/flexicm/layers/layers.py b/flexicm/layers/layers.py new file mode 100644 index 0000000000000000000000000000000000000000..73712292427893af0c3873f3a06bba8711b25e51 --- /dev/null +++ b/flexicm/layers/layers.py @@ -0,0 +1,769 @@ +# Copyright (c) 2021-2022, InterDigital Communications, Inc +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted (subject to the limitations in the disclaimer +# below) provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# * Neither the name of InterDigital Communications, Inc nor the names of its +# contributors may be used to endorse or promote products derived from this +# software without specific prior written permission. + +# NO EXPRESS OR IMPLIED LICENSES TO ANY PARTY'S PATENT RIGHTS ARE GRANTED BY +# THIS LICENSE. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND +# CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT +# NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A +# PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR +# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, +# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, +# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; +# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, +# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR +# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF +# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.autograd import Function +import torch.utils.checkpoint as checkpoint + +from functools import reduce +from operator import mul +import math + +from timm.models.layers import DropPath, to_2tuple, trunc_normal_ + +from .gdn import GDN + +__all__ = [ + "AttentionBlock", + "MaskedConv2d", + "MultistageMaskedConv2d", + "ResidualBlock", + "ResidualBlockUpsample", + "ResidualBlockWithStride", + "conv3x3", + "subpel_conv3x3", + "QReLU", + "RSTB" +] + + +class MaskedConv2d(nn.Conv2d): + r"""Masked 2D convolution implementation, mask future "unseen" pixels. + Useful for building auto-regressive network components. + + Introduced in `"Conditional Image Generation with PixelCNN Decoders" + `_. + + Inherits the same arguments as a `nn.Conv2d`. Use `mask_type='A'` for the + first layer (which also masks the "current pixel"), `mask_type='B'` for the + following layers. + """ + + def __init__(self, *args: Any, mask_type: str = "A", **kwargs: Any): + super().__init__(*args, **kwargs) + + if mask_type not in ("A", "B"): + raise ValueError(f'Invalid "mask_type" value "{mask_type}"') + + self.register_buffer("mask", torch.ones_like(self.weight.data)) + _, _, h, w = self.mask.size() + self.mask[:, :, h // 2, w // 2 + (mask_type == "B") :] = 0 + self.mask[:, :, h // 2 + 1 :] = 0 + + def forward(self, x: Tensor) -> Tensor: + # TODO(begaintj): weight assigment is not supported by torchscript + self.weight.data *= self.mask + return super().forward(x) + + +class MultistageMaskedConv2d(nn.Conv2d): + def __init__(self, *args: Any, mask_type: str = "A", **kwargs: Any): + super().__init__(*args, **kwargs) + + self.register_buffer("mask", torch.zeros_like(self.weight.data)) + + if mask_type == 'A': + self.mask[:, :, 0::2, 0::2] = 1 + elif mask_type == 'B': + self.mask[:, :, 0::2, 1::2] = 1 + self.mask[:, :, 1::2, 0::2] = 1 + elif mask_type == 'C': + self.mask[:, :, :, :] = 1 + self.mask[:, :, 1:2, 1:2] = 0 + else: + raise ValueError(f'Invalid "mask_type" value "{mask_type}"') + + def forward(self, x: Tensor) -> Tensor: + # TODO: weight assigment is not supported by torchscript + self.weight.data *= self.mask + return super().forward(x) + + +def conv3x3(in_ch: int, out_ch: int, stride: int = 1) -> nn.Module: + """3x3 convolution with padding.""" + return nn.Conv2d(in_ch, out_ch, kernel_size=3, stride=stride, padding=1) + + +def subpel_conv3x3(in_ch: int, out_ch: int, r: int = 1) -> nn.Sequential: + """3x3 sub-pixel convolution for up-sampling.""" + return nn.Sequential( + nn.Conv2d(in_ch, out_ch * r**2, kernel_size=3, padding=1), nn.PixelShuffle(r) + ) + + +def conv1x1(in_ch: int, out_ch: int, stride: int = 1) -> nn.Module: + """1x1 convolution.""" + return nn.Conv2d(in_ch, out_ch, kernel_size=1, stride=stride) + + +class ResidualBlockWithStride(nn.Module): + """Residual block with a stride on the first convolution. + + Args: + in_ch (int): number of input channels + out_ch (int): number of output channels + stride (int): stride value (default: 2) + """ + + def __init__(self, in_ch: int, out_ch: int, stride: int = 2): + super().__init__() + self.conv1 = conv3x3(in_ch, out_ch, stride=stride) + self.leaky_relu = nn.LeakyReLU(inplace=True) + self.conv2 = conv3x3(out_ch, out_ch) + self.gdn = GDN(out_ch) + if stride != 1 or in_ch != out_ch: + self.skip = conv1x1(in_ch, out_ch, stride=stride) + else: + self.skip = None + + def forward(self, x: Tensor) -> Tensor: + identity = x + out = self.conv1(x) + out = self.leaky_relu(out) + out = self.conv2(out) + out = self.gdn(out) + + if self.skip is not None: + identity = self.skip(x) + + out += identity + return out + + +class ResidualBlockUpsample(nn.Module): + """Residual block with sub-pixel upsampling on the last convolution. + + Args: + in_ch (int): number of input channels + out_ch (int): number of output channels + upsample (int): upsampling factor (default: 2) + """ + + def __init__(self, in_ch: int, out_ch: int, upsample: int = 2): + super().__init__() + self.subpel_conv = subpel_conv3x3(in_ch, out_ch, upsample) + self.leaky_relu = nn.LeakyReLU(inplace=True) + self.conv = conv3x3(out_ch, out_ch) + self.igdn = GDN(out_ch, inverse=True) + self.upsample = subpel_conv3x3(in_ch, out_ch, upsample) + + def forward(self, x: Tensor) -> Tensor: + identity = x + out = self.subpel_conv(x) + out = self.leaky_relu(out) + out = self.conv(out) + out = self.igdn(out) + identity = self.upsample(x) + out += identity + return out + + +class ResidualBlock(nn.Module): + """Simple residual block with two 3x3 convolutions. + + Args: + in_ch (int): number of input channels + out_ch (int): number of output channels + """ + + def __init__(self, in_ch: int, out_ch: int): + super().__init__() + self.conv1 = conv3x3(in_ch, out_ch) + self.leaky_relu = nn.LeakyReLU(inplace=True) + self.conv2 = conv3x3(out_ch, out_ch) + if in_ch != out_ch: + self.skip = conv1x1(in_ch, out_ch) + else: + self.skip = None + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.leaky_relu(out) + out = self.conv2(out) + out = self.leaky_relu(out) + + if self.skip is not None: + identity = self.skip(x) + + out = out + identity + return out + + +class AttentionBlock(nn.Module): + """Self attention block. + + Simplified variant from `"Learned Image Compression with + Discretized Gaussian Mixture Likelihoods and Attention Modules" + `_, by Zhengxue Cheng, Heming Sun, Masaru + Takeuchi, Jiro Katto. + + Args: + N (int): Number of channels) + """ + + def __init__(self, N: int): + super().__init__() + + class ResidualUnit(nn.Module): + """Simple residual unit.""" + + def __init__(self): + super().__init__() + self.conv = nn.Sequential( + conv1x1(N, N // 2), + nn.ReLU(inplace=True), + conv3x3(N // 2, N // 2), + nn.ReLU(inplace=True), + conv1x1(N // 2, N), + ) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x: Tensor) -> Tensor: + identity = x + out = self.conv(x) + out += identity + out = self.relu(out) + return out + + self.conv_a = nn.Sequential(ResidualUnit(), ResidualUnit(), ResidualUnit()) + + self.conv_b = nn.Sequential( + ResidualUnit(), + ResidualUnit(), + ResidualUnit(), + conv1x1(N, N), + ) + + def forward(self, x: Tensor) -> Tensor: + identity = x + a = self.conv_a(x) + b = self.conv_b(x) + out = a * torch.sigmoid(b) + out += identity + return out + + +class QReLU(Function): + """QReLU + + Clamping input with given bit-depth range. + Suppose that input data presents integer through an integer network + otherwise any precision of input will simply clamp without rounding + operation. + + Pre-computed scale with gamma function is used for backward computation. + + More details can be found in + `"Integer networks for data compression with latent-variable models" + `_, + by Johannes Ballé, Nick Johnston and David Minnen, ICLR in 2019 + + Args: + input: a tensor data + bit_depth: source bit-depth (used for clamping) + beta: a parameter for modeling the gradient during backward computation + """ + + @staticmethod + def forward(ctx, input, bit_depth, beta): + # TODO(choih): allow to use adaptive scale instead of + # pre-computed scale with gamma function + ctx.alpha = 0.9943258522851727 + ctx.beta = beta + ctx.max_value = 2**bit_depth - 1 + ctx.save_for_backward(input) + + return input.clamp(min=0, max=ctx.max_value) + + @staticmethod + def backward(ctx, grad_output): + grad_input = None + (input,) = ctx.saved_tensors + + grad_input = grad_output.clone() + grad_sub = ( + torch.exp( + (-ctx.alpha**ctx.beta) + * torch.abs(2.0 * input / ctx.max_value - 1) ** ctx.beta + ) + * grad_output.clone() + ) + + grad_input[input < 0] = grad_sub[input < 0] + grad_input[input > ctx.max_value] = grad_sub[input > ctx.max_value] + + return grad_input, None, None + + +class PatchEmbed(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + x = x.flatten(2).transpose(1, 2) # B Ph*Pw C + return x + + def flops(self): + flops = 0 + return flops + + +class PatchUnEmbed(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x, x_size): + B, HW, C = x.shape + x = x.transpose(1, 2).view(B, -1, x_size[0], x_size[1]) + return x + + def flops(self): + flops = 0 + return flops + + +class Mlp(nn.Module): + def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x): + x = self.fc1(x) + x = self.act(x) + x = self.drop(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +def window_partition(x, window_size): + """ + Args: + x: (B, H, W, C) + window_size (int): window size + Returns: + windows: (num_windows*B, window_size, window_size, C) + """ + B, H, W, C = x.shape + x = x.view(B, H // window_size, window_size, W // window_size, window_size, C) + windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C) + return windows + + +def window_reverse(windows, window_size, H, W): + """ + Args: + windows: (num_windows*B, window_size, window_size, C) + window_size (int): Window size + H (int): Height of image + W (int): Width of image + Returns: + x: (B, H, W, C) + """ + B = int(windows.shape[0] / (H * W / window_size / window_size)) + x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1) + x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1) + return x + + +class WindowAttention(nn.Module): + r""" Window based multi-head self attention (W-MSA) module with relative position bias. + It supports both of shifted and non-shifted window. + Args: + dim (int): Number of input channels. + window_size (tuple[int]): The height and width of the window. + num_heads (int): Number of attention heads. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set + attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 + proj_drop (float, optional): Dropout ratio of output. Default: 0.0 + """ + + def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.): + + super().__init__() + self.dim = dim + self.window_size = window_size # Wh, Ww + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = qk_scale or head_dim ** -0.5 + + # define a parameter table of relative position bias + self.relative_position_bias_table = nn.Parameter( + torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH + + # get pair-wise relative position index for each token inside the window + coords_h = torch.arange(self.window_size[0]) + coords_w = torch.arange(self.window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww + coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww + relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2 + relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0 + relative_coords[:, :, 1] += self.window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 + relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww + self.register_buffer("relative_position_index", relative_position_index) + + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = nn.Dropout(attn_drop) + self.proj = nn.Linear(dim, dim) + + self.proj_drop = nn.Dropout(proj_drop) + + trunc_normal_(self.relative_position_bias_table, std=.02) + self.softmax = nn.Softmax(dim=-1) + + def forward(self, x, mask=None): + """ + Args: + x: input features with shape of (num_windows*B, N, C) + mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None + """ + out_vis = dict() + B_, N, C = x.shape + qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple) + + q = q * self.scale + attn = (q @ k.transpose(-2, -1)) + + out_vis['inner_prod'] = attn.detach() + + relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view( + self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww + + out_vis['rpb'] = relative_position_bias.unsqueeze(0).detach() + + attn = attn + relative_position_bias.unsqueeze(0) + + if mask is not None: + nW = mask.shape[0] + attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, self.num_heads, N, N) + attn = self.softmax(attn) + else: + attn = self.softmax(attn) + + out_vis['attn'] = attn.detach() + + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(1, 2).reshape(B_, N, C) + x = self.proj(x) + x = self.proj_drop(x) + return x, out_vis + + def extra_repr(self) -> str: + return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}' + + def flops(self, N, img_N): + # calculate flops for 1 window with token length of N + flops = 0 + # qkv = self.qkv(x) + flops += N * self.dim * 3 * self.dim + # attn = (q @ k.transpose(-2, -1)) + flops += self.num_heads * img_N * (self.dim // self.num_heads) * N + # x = (attn @ v) + flops += self.num_heads * img_N * N * (self.dim // self.num_heads) + # x = self.proj(x) + flops += img_N * self.dim * self.dim + return flops + + + +class SwinTransformerBlock(nn.Module): + r""" Swin Transformer Block. + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resulotion. + num_heads (int): Number of attention heads. + window_size (int): Window size. + shift_size (int): Shift size for SW-MSA. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float, optional): Stochastic depth rate. Default: 0.0 + act_layer (nn.Module, optional): Activation layer. Default: nn.GELU + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + """ + + def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0., + act_layer=nn.GELU, norm_layer=nn.LayerNorm): + super().__init__() + self.dim = dim + self.input_resolution = input_resolution + self.actual_resolution = None + self.num_heads = num_heads + self.window_size = window_size + self.shift_size = shift_size + self.mlp_ratio = mlp_ratio + if min(self.input_resolution) <= self.window_size: + # if window size is larger than input resolution, we don't partition windows + self.shift_size = 0 + self.window_size = min(self.input_resolution) + assert 0 <= self.shift_size < self.window_size, "shift_size must in 0-window_size" + + self.norm1 = norm_layer(dim) + self.attn = WindowAttention( + dim, window_size=to_2tuple(self.window_size), num_heads=num_heads, + qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop) + + self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() + self.norm2 = norm_layer(dim) + mlp_hidden_dim = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) + + if self.shift_size > 0: + attn_mask = None #self.calculate_mask(self.input_resolution) + else: + attn_mask = None + + self.register_buffer("attn_mask", attn_mask) + + def calculate_mask(self, x_size): + # calculate attention mask for SW-MSA + H, W = x_size + img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1 + h_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + w_slices = (slice(0, -self.window_size), + slice(-self.window_size, -self.shift_size), + slice(-self.shift_size, None)) + cnt = 0 + for h in h_slices: + for w in w_slices: + img_mask[:, h, w, :] = cnt + cnt += 1 + + mask_windows = window_partition(img_mask, self.window_size) # nW, window_size, window_size, 1 + mask_windows = mask_windows.view(-1, self.window_size * self.window_size) + attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) + attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0)) + + return attn_mask + + def forward(self, x, x_size): + self.actual_resolution = x_size + H, W = x_size + B, L, C = x.shape + # assert L == H * W, "input feature has wrong size" + + shortcut = x + x = self.norm1(x) + x = x.view(B, H, W, C) + + # cyclic shift + if self.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2)) + else: + shifted_x = x + + # partition windows + x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C + x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C + + # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size + if self.input_resolution == x_size: + attn_windows, out_vis = self.attn(x_windows, mask=None) # nW*B, window_size*window_size, C + else: + attn_windows, out_vis = self.attn(x_windows, mask=None) + + # merge windows + attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C) + shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C + + # reverse cyclic shift + if self.shift_size > 0: + x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2)) + else: + x = shifted_x + x = x.view(B, H * W, C) + + # FFN + x = shortcut + self.drop_path(x) + x = x + self.drop_path(self.mlp(self.norm2(x))) + + return x, out_vis + + def extra_repr(self) -> str: + return f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, " \ + f"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}" + + def flops(self): + flops = 0 + H, W = self.input_resolution + # norm1 + flops += self.dim * H * W + # W-MSA/SW-MSA + nW = H * W / self.window_size / self.window_size + flops += nW * self.attn.flops(self.window_size * self.window_size) + # mlp + flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio + # norm2 + flops += self.dim * H * W + return flops + +class BasicLayer(nn.Module): + """ A basic Swin Transformer layer for one stage. + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + window_size (int): Local window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + + def __init__(self, dim, input_resolution, depth, num_heads, window_size, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., norm_layer=nn.LayerNorm, use_checkpoint=False, + block_module=SwinTransformerBlock): + + super().__init__() + self.dim = dim + self.input_resolution = input_resolution + self.depth = depth + self.use_checkpoint = use_checkpoint + + + self.blocks = nn.ModuleList([ + block_module( + dim=dim, input_resolution=input_resolution, + num_heads=num_heads, window_size=window_size, + shift_size=0 if (i % 2 == 0) else window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop, attn_drop=attn_drop, + drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path, # noqa + norm_layer=norm_layer) + for i in range(depth)]) + + + + def forward(self, x, x_size): + attns = [] + for i, blk in enumerate(self.blocks): + if self.use_checkpoint: + x = checkpoint.checkpoint(blk, x) + else: + x, _ = blk(x, x_size) + attn = None + attns.append(attn) + + return x, attns + + def extra_repr(self) -> str: + return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}" + + def flops(self): + flops = 0 + for blk in self.blocks: + flops += blk.flops() + return flops + + +class RSTB(nn.Module): + """Residual Swin Transformer Block (RSTB). + Args: + dim (int): Number of input channels. + input_resolution (tuple[int]): Input resolution. + depth (int): Number of blocks. + num_heads (int): Number of attention heads. + window_size (int): Local window size. + mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. + qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True + qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. + drop (float, optional): Dropout rate. Default: 0.0 + attn_drop (float, optional): Attention dropout rate. Default: 0.0 + drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0 + norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm + use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False. + """ + def __init__(self, dim, input_resolution, depth, num_heads, window_size, + mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., + drop_path=0., norm_layer=nn.LayerNorm, use_checkpoint=False): + super(RSTB, self).__init__() + + self.dim = dim + self.input_resolution = input_resolution + + self.residual_group = BasicLayer(dim=dim, + input_resolution=input_resolution, + depth=depth, + num_heads=num_heads, + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, qk_scale=qk_scale, + drop=drop, attn_drop=attn_drop, + drop_path=drop_path, + norm_layer=norm_layer, + use_checkpoint=use_checkpoint + ) + + self.patch_embed = PatchEmbed() + self.patch_unembed = PatchUnEmbed() + + + def forward(self, x, x_size): + out = self.patch_embed(x) + out, attns = self.residual_group(out, x_size) + return self.patch_unembed(out, x_size) + x, attns + + def flops(self): + flops = 0 + flops += self.residual_group.flops() + flops += self.patch_embed.flops() + flops += self.patch_unembed.flops() + + return flops diff --git a/flexicm/models/__init__.py b/flexicm/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..47768a532db1c0f7b1c6d36877f1708832216d2a --- /dev/null +++ b/flexicm/models/__init__.py @@ -0,0 +1,14 @@ +from .sfma import SFMA +from .task_connector import TaskConnector +from .taic import TAIC +from .ctaic import CTAIC +from .conditional import ConditionalPromptGenerator, ConditionGenerator + +__all__ = [ + "SFMA", + "TaskConnector", + "TAIC", + "CTAIC", + "ConditionalPromptGenerator", + "ConditionGenerator", +] diff --git a/flexicm/models/conditional.py b/flexicm/models/conditional.py new file mode 100644 index 0000000000000000000000000000000000000000..9000811c33a65d907efff6cfa8f7cdd96568206e --- /dev/null +++ b/flexicm/models/conditional.py @@ -0,0 +1,140 @@ +"""Conditional Prompt Generator and Condition Generator for C-TAIC (Fig. 2).""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from compressai.models.utils import conv, deconv + + +def window_partition_features(x, window_size): + """Partition (B,C,H,W) into (B*nW, window_size*window_size, C).""" + B, C, H, W = x.shape + x = x.view(B, C, H // window_size, window_size, W // window_size, window_size) + x = x.permute(0, 2, 4, 3, 5, 1).contiguous() + return x.view(-1, window_size * window_size, C) + + +class MaskGenerator(nn.Module): + """Lightweight soft mask over fused base-latent / image features.""" + + def __init__(self, img_channels=3, latent_channels=192, out_channels=192, mid=64): + super().__init__() + self.img_stem = nn.Sequential( + conv(img_channels, mid, kernel_size=5, stride=2), # H/2 + nn.GELU(), + conv(mid, mid, kernel_size=3, stride=2), # H/4 + nn.GELU(), + conv(mid, mid, kernel_size=3, stride=2), # H/8 + nn.GELU(), + conv(mid, out_channels, kernel_size=3, stride=2), # H/16 + ) + self.latent_proj = nn.Conv2d(latent_channels, out_channels, 1) + self.fuse = nn.Sequential( + nn.Conv2d(out_channels * 2, mid, 1), + nn.GELU(), + nn.Conv2d(mid, out_channels, 1), + nn.Sigmoid(), + ) + + def forward(self, x, y_b_hat): + fx = self.img_stem(x) + fy = self.latent_proj(y_b_hat) + if fx.shape[-2:] != fy.shape[-2:]: + fx = F.interpolate(fx, size=fy.shape[-2:], mode="bilinear", align_corners=False) + m = self.fuse(torch.cat([fx, fy], dim=1)) + return m, fx + + +class ConditionalPromptGenerator(nn.Module): + """Generate multi-scale prompts C2 / C4 from (x, y_b_hat). + + Prompts are returned as lists of per-window token tensors compatible with + TIC window_size=8 attention: + C2 -> first encoder STB at H/2 (16 prompt tokens / window) + C4 -> second encoder STB at H/4 (16 prompt tokens / window) + """ + + def __init__(self, latent_channels=192, prompt_dim=128, window_size=8): + super().__init__() + self.window_size = window_size + self.prompt_dim = prompt_dim + self.mask_gen = MaskGenerator( + img_channels=3, latent_channels=latent_channels, out_channels=latent_channels + ) + self.img_to_latent = nn.Sequential( + conv(3, 64, kernel_size=5, stride=2), + nn.GELU(), + conv(64, 128, kernel_size=3, stride=2), + nn.GELU(), + conv(128, latent_channels, kernel_size=3, stride=2), + nn.GELU(), + conv(latent_channels, latent_channels, kernel_size=3, stride=2), + ) + # H/16 -> H/8 -> H/4 + self.up1 = deconv(latent_channels, latent_channels, kernel_size=3, stride=2) + self.up2 = deconv(latent_channels, latent_channels, kernel_size=3, stride=2) + self.proj_c4 = nn.Conv2d(latent_channels, prompt_dim, 1) + self.proj_c2 = nn.Conv2d(latent_channels, prompt_dim, 1) + + def _to_window_prompts(self, feat, target_hw): + """Map spatial prompt feature to windows of the target STB resolution. + + Each TIC window (window_size x window_size) at target_hw covers a + (window_size/scale) block on ``feat``, yielding 16 tokens when + feat is 2x down relative to target (paper ratio 1/4 of 64). + """ + B, C, Hf, Wf = feat.shape + Ht, Wt = target_hw + # Align spatial size: target windows expect feat at Ht/2 x Wt/2 + expect_h, expect_w = Ht // 2, Wt // 2 + if (Hf, Wf) != (expect_h, expect_w): + feat = F.interpolate(feat, size=(expect_h, expect_w), mode="bilinear", align_corners=False) + # Partition with window_size//2 so each target window gets 16 tokens + ws = self.window_size // 2 + return window_partition_features(feat, ws) + + def forward(self, x, y_b_hat, sizes_h2, sizes_h4): + """ + Args: + x: input image (B,3,H,W) + y_b_hat: base-layer latent after AD (B,192,H/16,W/16) + sizes_h2: (H/2, W/2) of first STB + sizes_h4: (H/4, W/4) of second STB + Returns: + prompt_c2, prompt_c4: (B*nW, 16, prompt_dim) + """ + m, _ = self.mask_gen(x, y_b_hat) + fx = self.img_to_latent(x) + if fx.shape[-2:] != y_b_hat.shape[-2:]: + fx = F.interpolate(fx, size=y_b_hat.shape[-2:], mode="bilinear", align_corners=False) + if m.shape[-2:] != y_b_hat.shape[-2:]: + m = F.interpolate(m, size=y_b_hat.shape[-2:], mode="bilinear", align_corners=False) + f_sum = m * fx + (1.0 - m) * y_b_hat + + f_h8 = self.up1(f_sum) # H/8 + f_h4 = self.up2(f_h8) # H/4 + + c4_map = self.proj_c4(f_h8) + c2_map = self.proj_c2(f_h4) + + prompt_c2 = self._to_window_prompts(c2_map, sizes_h2) + prompt_c4 = self._to_window_prompts(c4_map, sizes_h4) + return prompt_c2, prompt_c4 + + +class ConditionGenerator(nn.Module): + """Decoder-side condition Cd from base latent y_b_hat. + + Lightweight: Deconv (H/16->H/8) + 1x1 Linear/Conv to match Task Connector channels. + """ + + def __init__(self, latent_channels=192, out_channels=128): + super().__init__() + self.net = nn.Sequential( + deconv(latent_channels, out_channels, kernel_size=3, stride=2), + nn.GELU(), + nn.Conv2d(out_channels, out_channels, kernel_size=1), + ) + + def forward(self, y_b_hat): + return self.net(y_b_hat) diff --git a/flexicm/models/cross_attention.py b/flexicm/models/cross_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..3abd52277418291920775425dc81cb2e005fb3ad --- /dev/null +++ b/flexicm/models/cross_attention.py @@ -0,0 +1,118 @@ +"""Cross-attention with conditional prompts inside TIC window attention. + +Q is computed from feature tokens only; K,V from [features; prompts] (paper Eq.5). +Uses the frozen TIC qkv projection weights. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from flexicm.layers.layers import window_partition, window_reverse + + +def window_attention_with_prompt(attn_module, x_windows, prompts, mask=None): + """ + Args: + attn_module: WindowAttention with attributes qkv, proj, scale, num_heads, ... + x_windows: (B*nW, N, C) feature tokens in each window + prompts: (B*nW, Np, C) prompt tokens aligned with windows + mask: optional attention mask for shifted windows (N x N); prompt cols get 0 + Returns: + out: (B*nW, N, C) + """ + B_, N, C = x_windows.shape + Np = prompts.shape[1] + num_heads = attn_module.num_heads + head_dim = C // num_heads + scale = attn_module.scale + + # Q from features only + qkv_f = attn_module.qkv(x_windows).reshape(B_, N, 3, num_heads, head_dim).permute(2, 0, 3, 1, 4) + q, k_f, v_f = qkv_f[0], qkv_f[1], qkv_f[2] + + # K,V from prompts + qkv_p = attn_module.qkv(prompts).reshape(B_, Np, 3, num_heads, head_dim).permute(2, 0, 3, 1, 4) + k_p, v_p = qkv_p[1], qkv_p[2] + + k = torch.cat([k_f, k_p], dim=2) # (B_, heads, N+Np, head_dim) + v = torch.cat([v_f, v_p], dim=2) + + q = q * scale + attn = q @ k.transpose(-2, -1) # (B_, heads, N, N+Np) + + # Relative position bias only on feature-feature block + relative_position_bias = attn_module.relative_position_bias_table[ + attn_module.relative_position_index.view(-1) + ].view( + attn_module.window_size[0] * attn_module.window_size[1], + attn_module.window_size[0] * attn_module.window_size[1], + -1, + ) + relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() + attn[:, :, :, :N] = attn[:, :, :, :N] + relative_position_bias.unsqueeze(0) + + if mask is not None: + # mask: (nW, N, N) -> pad prompt columns with 0 + nW = mask.shape[0] + mask_pad = F.pad(mask, (0, Np), value=0.0) # (nW, N, N+Np) + attn = attn.view(-1, nW, num_heads, N, N + Np) + mask_pad.unsqueeze(1).unsqueeze(0) + attn = attn.view(-1, num_heads, N, N + Np) + + attn = attn_module.softmax(attn) + attn = attn_module.attn_drop(attn) + out = (attn @ v).transpose(1, 2).reshape(B_, N, C) + out = attn_module.proj(out) + out = attn_module.proj_drop(out) + return out + + +def swin_block_forward_with_prompt(block, x, x_size, prompts): + """Run one SwinTransformerBlock with optional prompt cross-attention. + + Args: + block: SwinTransformerBlock + x: (B, H*W, C) + x_size: (H, W) + prompts: (B*nW, Np, C) or None + """ + H, W = x_size + B, L, C = x.shape + shortcut = x + x = block.norm1(x) + x = x.view(B, H, W, C) + + if block.shift_size > 0: + shifted_x = torch.roll(x, shifts=(-block.shift_size, -block.shift_size), dims=(1, 2)) + attn_mask = block.calculate_mask(x_size).to(x.device) + else: + shifted_x = x + attn_mask = None + + x_windows = window_partition(shifted_x, block.window_size) + x_windows = x_windows.view(-1, block.window_size * block.window_size, C) + + if prompts is not None: + attn_windows = window_attention_with_prompt(block.attn, x_windows, prompts, mask=attn_mask) + else: + attn_windows, _ = block.attn(x_windows, mask=attn_mask) + + attn_windows = attn_windows.view(-1, block.window_size, block.window_size, C) + shifted_x = window_reverse(attn_windows, block.window_size, H, W) + + if block.shift_size > 0: + x = torch.roll(shifted_x, shifts=(block.shift_size, block.shift_size), dims=(1, 2)) + else: + x = shifted_x + x = x.view(B, H * W, C) + x = shortcut + block.drop_path(x) + x = x + block.drop_path(block.mlp(block.norm2(x))) + return x + + +def rstb_forward_with_prompt(rstb, x, x_size, prompts=None): + """RSTB forward; if prompts is set, inject into every block of the RSTB.""" + out = rstb.patch_embed(x) + for blk in rstb.residual_group.blocks: + out = swin_block_forward_with_prompt(blk, out, x_size, prompts) + return rstb.patch_unembed(out, x_size) + x diff --git a/flexicm/models/ctaic.py b/flexicm/models/ctaic.py new file mode 100644 index 0000000000000000000000000000000000000000..9f2a6874e19f9375c7a73db6cbe109b3bf1556a0 --- /dev/null +++ b/flexicm/models/ctaic.py @@ -0,0 +1,100 @@ +"""C-TAIC: Conditional Task-Adaptive Image Coding (FlexICM extension layer, Fig. 2). + +Stage-1 (TAIC mode): train SFMA + Task Connector (no base-layer condition). +Stage-2 (C-TAIC mode): freeze all except Conditional Prompt Generator + Condition Generator; + use base-layer y_b_hat for cross-attention prompts and decoder Cd. +""" + +import torch +import torch.nn as nn + +from flexicm.models.taic import TAIC +from flexicm.models.conditional import ConditionalPromptGenerator, ConditionGenerator + + +class CTAIC(TAIC): + def __init__( + self, + N=128, + M=192, + input_resolution=(256, 256), + out_channels=128, + in_channel=3, + ): + super().__init__( + N=N, + M=M, + input_resolution=input_resolution, + out_channels=out_channels, + in_channel=in_channel, + ) + self.prompt_generator = ConditionalPromptGenerator( + latent_channels=M, prompt_dim=N, window_size=8 + ) + self.condition_generator = ConditionGenerator( + latent_channels=M, out_channels=N + ) + + def forward(self, x, y_b_hat=None, use_condition=True): + """ + Args: + x: input image + y_b_hat: base-layer latent after AD (H/16 x W/16 x M). Required when use_condition. + use_condition: if False, operate as TAIC (stage-1 / graceful degradation) + """ + prompts = None + condition = None + if use_condition and y_b_hat is not None: + H, W = x.shape[2], x.shape[3] + c2, c4 = self.prompt_generator( + x, y_b_hat, sizes_h2=(H // 2, W // 2), sizes_h4=(H // 4, W // 4) + ) + prompts = {"c2": c2, "c4": c4} + condition = self.condition_generator(y_b_hat) + return super().forward(x, prompts=prompts, condition=condition) + + def compress(self, x, y_b_hat=None, use_condition=True): + prompts = None + if use_condition and y_b_hat is not None: + H, W = x.shape[2], x.shape[3] + c2, c4 = self.prompt_generator( + x, y_b_hat, sizes_h2=(H // 2, W // 2), sizes_h4=(H // 4, W // 4) + ) + prompts = {"c2": c2, "c4": c4} + return super().compress(x, prompts=prompts) + + def decompress(self, strings, shape, x_size=None, y_b_hat=None, use_condition=True): + condition = None + if use_condition and y_b_hat is not None: + condition = self.condition_generator(y_b_hat) + return super().decompress(strings, shape, x_size=x_size, condition=condition) + + def freeze_for_stage1(self): + """Train SFMA + Task Connector only (TAIC mode).""" + for name, p in self.named_parameters(): + train = ("encoder_sfmas" in name) or ("task_connector" in name) + p.requires_grad = train + # keep prompt/condition gens frozen in stage-1 + for p in self.prompt_generator.parameters(): + p.requires_grad = False + for p in self.condition_generator.parameters(): + p.requires_grad = False + + def freeze_for_stage2(self): + """Train Conditional Prompt Generator + Condition Generator only.""" + for p in self.parameters(): + p.requires_grad = False + for p in self.prompt_generator.parameters(): + p.requires_grad = True + for p in self.condition_generator.parameters(): + p.requires_grad = True + + def load_taic_checkpoint(self, state_dict, strict=False): + """Initialize from a trained TAIC (extension-task) checkpoint.""" + own = self.state_dict() + filtered = {} + for k, v in state_dict.items(): + nk = k[7:] if k.startswith("module.") else k + if nk in own and own[nk].shape == v.shape: + filtered[nk] = v + return self.load_state_dict(filtered, strict=False) diff --git a/flexicm/models/sfma.py b/flexicm/models/sfma.py new file mode 100644 index 0000000000000000000000000000000000000000..b6c17027218d5c385e31b064c0dd4ab3e2713549 --- /dev/null +++ b/flexicm/models/sfma.py @@ -0,0 +1,40 @@ +"""Spatial-Frequency Modulation Adapter (SFMA), same design as AdaptiveICMH.""" + +import torch +import torch.nn as nn + + +class SFMA(nn.Module): + def __init__(self, in_dim=128, middle_dim=64, adapt_factor=1.0): + super().__init__() + self.factor = adapt_factor + # Spatial branch + self.s_down1 = nn.Conv2d(in_dim, middle_dim, 1, 1, 0) + self.s_down2 = nn.Conv2d(in_dim, middle_dim, 1, 1, 0) + self.s_dw = nn.Conv2d(middle_dim, middle_dim, 5, 1, 2, groups=middle_dim) + self.s_relu = nn.ReLU(inplace=True) + self.s_up = nn.Conv2d(middle_dim, in_dim, 1, 1, 0) + # Frequency branch + self.f_down = nn.Conv2d(in_dim, middle_dim, 1, 1, 0) + self.f_relu1 = nn.ReLU(inplace=True) + self.f_relu2 = nn.ReLU(inplace=True) + self.f_up = nn.Conv2d(middle_dim, in_dim, 1, 1, 0) + self.f_dw = nn.Conv2d(middle_dim, middle_dim, 3, 1, 1, groups=middle_dim) + self.f_inter = nn.Conv2d(middle_dim, middle_dim, 1, 1, 0) + self.sg = nn.Sigmoid() + + def forward(self, x): + _, _, H, W = x.shape + y = torch.fft.rfft2(self.f_down(x), dim=(2, 3), norm="backward") + y_amp = torch.abs(y) + y_phs = torch.angle(y) + y_amp_modulation = self.f_inter(self.f_relu1(self.f_dw(y_amp))) + y_amp = y_amp * self.sg(y_amp_modulation) + y_real = y_amp * torch.cos(y_phs) + y_img = y_amp * torch.sin(y_phs) + y = torch.complex(y_real, y_img) + y = torch.fft.irfft2(y, s=(H, W), norm="backward") + + f_modulate = self.f_up(self.f_relu2(y)) + s_modulate = self.s_up(self.s_relu(self.s_dw(self.s_down1(x)) * self.s_down2(x))) + return x + (s_modulate + f_modulate) * self.factor diff --git a/flexicm/models/taic.py b/flexicm/models/taic.py new file mode 100644 index 0000000000000000000000000000000000000000..b30d534cb05fef3ac7b948b298a5990aa0a554ed --- /dev/null +++ b/flexicm/models/taic.py @@ -0,0 +1,395 @@ +"""TAIC: Task-Adaptive Image Coding (FlexICM base layer, Fig. 1). + +Encoder: frozen TIC + trainable SFMA (same placement as AdaptiveICMH). +Decoder: frozen TIC g_s0 (STB) + g_s1 (Deconv) + trainable Task Connector + -> h at H/4 x W/4 x out_channels (no full image reconstruction). +""" + +import math + +import torch +import torch.nn as nn +from compressai.entropy_models import EntropyBottleneck, GaussianConditional +from compressai.models.utils import conv, deconv, update_registered_buffers +from timm.models.layers import trunc_normal_ + +from flexicm.layers.layers import RSTB +from flexicm.models.sfma import SFMA +from flexicm.models.task_connector import TaskConnector + +SCALES_MIN = 0.11 +SCALES_MAX = 256 +SCALES_LEVELS = 64 + + +def ste_round(x): + return torch.round(x) - x.detach() + x + + +def get_scale_table(min=SCALES_MIN, max=SCALES_MAX, levels=SCALES_LEVELS): + return torch.exp(torch.linspace(math.log(min), math.log(max), levels)) + + +class TAIC(nn.Module): + def __init__( + self, + N=128, + M=192, + input_resolution=(256, 256), + out_channels=128, + in_channel=3, + ): + super().__init__() + depths = [2, 4, 6, 2, 2, 2] + num_heads = [8, 8, 8, 16, 16, 16] + window_size = 8 + mlp_ratio = 2.0 + qkv_bias = True + qk_scale = None + drop_rate = 0.0 + attn_drop_rate = 0.0 + drop_path_rate = 0.1 + norm_layer = nn.LayerNorm + use_checkpoint = False + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] + + self.N = N + self.M = M + self.out_channels = out_channels + + # Encoder-side SFMA (trainable) + self.encoder_sfmas = nn.Sequential(SFMA(N), SFMA(N), SFMA(N)) + + # ---- encoder (frozen TIC) ---- + self.g_a0 = conv(in_channel, N, kernel_size=5, stride=2) + self.g_a1 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 2, input_resolution[1] // 2), + depth=depths[0], + num_heads=num_heads[0], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:0]) : sum(depths[:1])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.g_a2 = conv(N, N, kernel_size=3, stride=2) + self.g_a3 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 4, input_resolution[1] // 4), + depth=depths[1], + num_heads=num_heads[1], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:1]) : sum(depths[:2])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.g_a4 = conv(N, N, kernel_size=3, stride=2) + self.g_a5 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 8, input_resolution[1] // 8), + depth=depths[2], + num_heads=num_heads[2], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:2]) : sum(depths[:3])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.g_a6 = conv(N, M, kernel_size=3, stride=2) + self.g_a7 = RSTB( + dim=M, + input_resolution=(input_resolution[0] // 16, input_resolution[1] // 16), + depth=depths[3], + num_heads=num_heads[3], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:3]) : sum(depths[:4])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + + # ---- hyperprior (frozen TIC) ---- + self.h_a0 = conv(M, N, kernel_size=3, stride=2) + self.h_a1 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 32, input_resolution[1] // 32), + depth=depths[4], + num_heads=num_heads[4], + window_size=window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:4]) : sum(depths[:5])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.h_a2 = conv(N, N, kernel_size=3, stride=2) + self.h_a3 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 64, input_resolution[1] // 64), + depth=depths[5], + num_heads=num_heads[5], + window_size=window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths[:5]) : sum(depths[:6])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + + depths_rev = depths[::-1] + num_heads_rev = num_heads[::-1] + self.h_s0 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 64, input_resolution[1] // 64), + depth=depths_rev[0], + num_heads=num_heads_rev[0], + window_size=window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths_rev[:0]) : sum(depths_rev[:1])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.h_s1 = deconv(N, N, kernel_size=3, stride=2) + self.h_s2 = RSTB( + dim=N, + input_resolution=(input_resolution[0] // 32, input_resolution[1] // 32), + depth=depths_rev[1], + num_heads=num_heads_rev[1], + window_size=window_size // 2, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths_rev[:1]) : sum(depths_rev[:2])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.h_s3 = deconv(N, M * 2, kernel_size=3, stride=2) + + self.entropy_bottleneck = EntropyBottleneck(N) + self.gaussian_conditional = GaussianConditional(None) + + # ---- partial decoder (frozen TIC g_s0 + g_s1) ---- + self.g_s0 = RSTB( + dim=M, + input_resolution=(input_resolution[0] // 16, input_resolution[1] // 16), + depth=depths_rev[2], + num_heads=num_heads_rev[2], + window_size=window_size, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[sum(depths_rev[:2]) : sum(depths_rev[:3])], + norm_layer=norm_layer, + use_checkpoint=use_checkpoint, + ) + self.g_s1 = deconv(M, N, kernel_size=3, stride=2) # H/16 -> H/8 + + # ---- trainable Task Connector ---- + self.task_connector = TaskConnector( + in_channels=N, mid_channels=N, out_channels=out_channels + ) + + self.apply(self._init_weights) + + def _init_weights(self, m): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + def g_a(self, x, x_size=None, prompts=None): + """ + Args: + prompts: optional dict {"c2": tensor, "c4": tensor} for C-TAIC + """ + if x_size is None: + x_size = x.shape[2:4] + x = self.g_a0(x) + if prompts is not None and prompts.get("c2") is not None: + from flexicm.models.cross_attention import rstb_forward_with_prompt + + x = rstb_forward_with_prompt( + self.g_a1, x, (x_size[0] // 2, x_size[1] // 2), prompts["c2"] + ) + else: + x, _ = self.g_a1(x, (x_size[0] // 2, x_size[1] // 2)) + x = self.encoder_sfmas[0](x) + x = self.g_a2(x) + + if prompts is not None and prompts.get("c4") is not None: + from flexicm.models.cross_attention import rstb_forward_with_prompt + + x = rstb_forward_with_prompt( + self.g_a3, x, (x_size[0] // 4, x_size[1] // 4), prompts["c4"] + ) + else: + x, _ = self.g_a3(x, (x_size[0] // 4, x_size[1] // 4)) + x = self.encoder_sfmas[1](x) + x = self.g_a4(x) + + x, _ = self.g_a5(x, (x_size[0] // 8, x_size[1] // 8)) + x = self.encoder_sfmas[2](x) + x = self.g_a6(x) + x, _ = self.g_a7(x, (x_size[0] // 16, x_size[1] // 16)) + return x + + def h_a(self, x, x_size=None): + if x_size is None: + x_size = (x.shape[2] * 16, x.shape[3] * 16) + x = self.h_a0(x) + x, _ = self.h_a1(x, (x_size[0] // 32, x_size[1] // 32)) + x = self.h_a2(x) + x, _ = self.h_a3(x, (x_size[0] // 64, x_size[1] // 64)) + return x + + def h_s(self, x, x_size=None): + if x_size is None: + x_size = (x.shape[2] * 64, x.shape[3] * 64) + x, _ = self.h_s0(x, (x_size[0] // 64, x_size[1] // 64)) + x = self.h_s1(x) + x, _ = self.h_s2(x, (x_size[0] // 32, x_size[1] // 32)) + x = self.h_s3(x) + return x + + def decode_feature(self, y_hat, x_size=None, condition=None): + """Map quantized latent to task feature h.""" + if x_size is None: + x_size = (y_hat.shape[2] * 16, y_hat.shape[3] * 16) + x, _ = self.g_s0(y_hat, (x_size[0] // 16, x_size[1] // 16)) + x = self.g_s1(x) # H/8 x W/8 x N + h = self.task_connector(x, condition=condition) + return h + + def aux_loss(self): + return sum(m.loss() for m in self.modules() if isinstance(m, EntropyBottleneck)) + + def forward(self, x, prompts=None, condition=None): + x_size = (x.shape[2], x.shape[3]) + y = self.g_a(x, x_size, prompts=prompts) + z = self.h_a(y, x_size) + _, z_likelihoods = self.entropy_bottleneck(z) + z_offset = self.entropy_bottleneck._get_medians() + z_hat = ste_round(z - z_offset) + z_offset + gaussian_params = self.h_s(z_hat, x_size) + scales_hat, means_hat = gaussian_params.chunk(2, 1) + _, y_likelihoods = self.gaussian_conditional(y, scales_hat, means=means_hat) + y_hat = ste_round(y - means_hat) + means_hat + h = self.decode_feature(y_hat, x_size, condition=condition) + return { + "h": h, + "y_hat": y_hat, + "likelihoods": {"y": y_likelihoods, "z": z_likelihoods}, + } + + def compress(self, x, prompts=None): + x_size = (x.shape[2], x.shape[3]) + y = self.g_a(x, x_size, prompts=prompts) + z = self.h_a(y, x_size) + z_strings = self.entropy_bottleneck.compress(z) + z_hat = self.entropy_bottleneck.decompress(z_strings, z.size()[-2:]) + gaussian_params = self.h_s(z_hat, x_size) + scales_hat, means_hat = gaussian_params.chunk(2, 1) + indexes = self.gaussian_conditional.build_indexes(scales_hat) + y_strings = self.gaussian_conditional.compress(y, indexes, means=means_hat) + return {"strings": [y_strings, z_strings], "shape": z.size()[-2:], "x_size": x_size} + + def decompress(self, strings, shape, x_size=None, condition=None): + z_hat = self.entropy_bottleneck.decompress(strings[1], shape) + if x_size is None: + x_size = (shape[0] * 64, shape[1] * 64) + gaussian_params = self.h_s(z_hat, x_size) + scales_hat, means_hat = gaussian_params.chunk(2, 1) + indexes = self.gaussian_conditional.build_indexes(scales_hat) + y_hat = self.gaussian_conditional.decompress(strings[0], indexes, means=means_hat) + h = self.decode_feature(y_hat, x_size, condition=condition) + return {"h": h, "y_hat": y_hat} + + def freeze_base_codec(self): + """Freeze TIC weights; keep SFMA + Task Connector trainable.""" + for name, p in self.named_parameters(): + if ("sfma" in name.lower()) or ("task_connector" in name): + p.requires_grad = True + else: + p.requires_grad = False + + def trainable_parameter_names(self): + return [n for n, p in self.named_parameters() if p.requires_grad] + + def update(self, scale_table=None, force=False): + if scale_table is None: + scale_table = get_scale_table() + self.gaussian_conditional.update_scale_table(scale_table, force=force) + updated = False + for m in self.children(): + if isinstance(m, EntropyBottleneck): + updated |= m.update(force=force) + return updated + + def load_state_dict(self, state_dict, strict=True): + update_registered_buffers( + self.entropy_bottleneck, + "entropy_bottleneck", + ["_quantized_cdf", "_offset", "_cdf_length"], + state_dict, + ) + update_registered_buffers( + self.gaussian_conditional, + "gaussian_conditional", + ["_quantized_cdf", "_offset", "_cdf_length", "scale_table"], + state_dict, + ) + super().load_state_dict(state_dict, strict=strict) + + def load_base_codec(self, state_dict, strict=False): + """Load pretrained TIC / TIC-SFMA weights into matching modules.""" + own = self.state_dict() + filtered = {} + for k, v in state_dict.items(): + nk = k[7:] if k.startswith("module.") else k + # skip decoder stages after g_s1 and any decoder SFMAs + if nk.startswith("g_s2") or nk.startswith("g_s3") or nk.startswith("g_s4"): + continue + if nk.startswith("g_s5") or nk.startswith("g_s6") or nk.startswith("g_s7"): + continue + if "decoder_sfmas" in nk: + continue + if nk in own and own[nk].shape == v.shape: + filtered[nk] = v + missing_unexpected = self.load_state_dict(filtered, strict=False) + return missing_unexpected diff --git a/flexicm/models/task_connector.py b/flexicm/models/task_connector.py new file mode 100644 index 0000000000000000000000000000000000000000..0528deec2e88e004ed094fd127639f7d383d6096 --- /dev/null +++ b/flexicm/models/task_connector.py @@ -0,0 +1,36 @@ +"""Lightweight Task Connector (Fig. 1 / Fig. 2). + +After frozen TIC decoder stages STB(g_s0)+Deconv(g_s1) at H/8 x W/8 x N, +the Task Connector produces h at H/4 x W/4 x out_channels via: + residual(Linear -> DW-Conv -> Linear) -> Deconv +""" + +import torch +import torch.nn as nn +from compressai.models.utils import deconv + + +class TaskConnector(nn.Module): + def __init__(self, in_channels=128, mid_channels=128, out_channels=128): + super().__init__() + self.linear1 = nn.Conv2d(in_channels, mid_channels, kernel_size=1) + self.dw = nn.Conv2d( + mid_channels, mid_channels, kernel_size=3, padding=1, groups=mid_channels + ) + self.act = nn.GELU() + self.linear2 = nn.Conv2d(mid_channels, in_channels, kernel_size=1) + # H/8 -> H/4 + self.upsample = deconv(in_channels, out_channels, kernel_size=3, stride=2) + + def forward(self, x, condition=None): + """ + Args: + x: (B, C, H/8, W/8) feature after frozen STB+Deconv + condition: optional (B, C, H/8, W/8) decoder-side condition Cd + """ + if condition is not None: + x = x + condition + residual = x + y = self.linear2(self.act(self.dw(self.linear1(x)))) + y = residual + y + return self.upsample(y) diff --git a/flexicm/tasks/__init__.py b/flexicm/tasks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..44ef136236e4200070d06e91b43c29a34213ee6f --- /dev/null +++ b/flexicm/tasks/__init__.py @@ -0,0 +1,219 @@ +"""Task-specific frozen teachers for FlexICM feature alignment. + +Five tasks (paper Sec.III.A / IV.A): + 1. Object detection - Faster R-CNN + Swin-B (FPN P2-P6) + 2. Instance segmentation - Mask R-CNN + Swin-B (FPN P2-P6) + 3. Semantic segmentation - UPerNet + Swin-B (FPN P2-P6) + 4. Panoptic segmentation - MaskFormer + Swin-B (stages F1-F4) + 5. Pose estimation - HigherHRNet (original HRNet backbone) + +Optional full MMDet/MMSeg/MMPose models can be attached for end-task evaluation; +training the codec only needs intermediate feature alignment. +""" + +from __future__ import annotations + +from typing import Dict, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from flexicm.tasks.losses import freeze_module +from flexicm.tasks.swin_teacher import SwinStageTeacher + + +class DetectionTeacher(nn.Module): + """Faster R-CNN / Mask R-CNN style: align FPN P2..P6.""" + + align_mode = "fpn" + out_channels = 128 # Swin-B F1 + + def __init__(self, pretrained_backbone: bool = True, task: str = "detection"): + super().__init__() + self.task = task + self.backbone = SwinStageTeacher(pretrained=pretrained_backbone, use_fpn=True) + freeze_module(self) + + def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]: + return self.backbone.gt_features(images) + + def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + return self.backbone.pred_features(h) + + +class SemanticSegTeacher(nn.Module): + """UPerNet style: align FPN features (paper Eq.2).""" + + align_mode = "fpn" + out_channels = 128 + + def __init__(self, pretrained_backbone: bool = True): + super().__init__() + self.backbone = SwinStageTeacher(pretrained=pretrained_backbone, use_fpn=True) + freeze_module(self) + + def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]: + return self.backbone.gt_features(images) + + def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + return self.backbone.pred_features(h) + + +class PanopticSegTeacher(nn.Module): + """MaskFormer style: align intermediate stages F1..F4 (paper Eq.3).""" + + align_mode = "stages" + out_channels = 128 + + def __init__(self, pretrained_backbone: bool = True): + super().__init__() + self.backbone = SwinStageTeacher(pretrained=pretrained_backbone, use_fpn=False) + freeze_module(self) + + def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]: + return self.backbone.gt_features(images) + + def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + return self.backbone.pred_features(h) + + +class HigherHRNetTeacher(nn.Module): + """Pose estimation teacher with original HRNet backbone (not Swin). + + Aligns multi-resolution HRNet features. h (H/4 x W/4 x C) is projected to + match the HRNet stem/stage-1 width, then remaining stages produce F1..F4-like maps. + """ + + align_mode = "stages" + out_channels = 32 # HRNet-W32 stem / stage channels (configurable) + + def __init__(self, width: int = 32, pretrained: bool = True): + super().__init__() + self.width = width + self.out_channels = width + self.stem = freeze_module(self._build_stem(width)) + self.stage_downs = freeze_module( + nn.ModuleList( + [ + nn.Sequential( + nn.Conv2d(width, width * 2, 3, stride=2, padding=1), + nn.BatchNorm2d(width * 2), + nn.ReLU(inplace=True), + ), + nn.Sequential( + nn.Conv2d(width * 2, width * 4, 3, stride=2, padding=1), + nn.BatchNorm2d(width * 4), + nn.ReLU(inplace=True), + ), + nn.Sequential( + nn.Conv2d(width * 4, width * 8, 3, stride=2, padding=1), + nn.BatchNorm2d(width * 8), + nn.ReLU(inplace=True), + ), + ] + ) + ) + self.h_proj = freeze_module(nn.Conv2d(128, width, 1)) + # Optional: load real HigherHRNet via mmpose if available + self.mmpose_model = None + if pretrained: + self._try_load_mmpose() + + @staticmethod + def _build_stem(width: int) -> nn.Module: + return nn.Sequential( + nn.Conv2d(3, 64, 3, stride=2, padding=1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, 64, 3, stride=2, padding=1), + nn.BatchNorm2d(64), + nn.ReLU(inplace=True), + nn.Conv2d(64, width, 3, padding=1), + nn.BatchNorm2d(width), + nn.ReLU(inplace=True), + ) + + def _try_load_mmpose(self): + try: + # Placeholder hook: users can replace with mmpose HigherHRNet + # e.g. init_pose_model(config, checkpoint) + self.mmpose_model = None + except Exception: + self.mmpose_model = None + + def _normalize(self, x: torch.Tensor) -> torch.Tensor: + mean = x.new_tensor([0.485, 0.456, 0.406])[None, :, None, None] + std = x.new_tensor([0.229, 0.224, 0.225])[None, :, None, None] + return (x - mean) / std + + def _stages_from_f1(self, f1: torch.Tensor) -> Dict[str, torch.Tensor]: + feats = [f1] + x = f1 + for down in self.stage_downs: + x = down(x) + feats.append(x) + return {f"f{i+1}": feats[i] for i in range(4)} + + def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]: + x = self._normalize(images) + f1 = self.stem(x) + return self._stages_from_f1(f1) + + def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + if h.shape[1] != self.width: + if h.shape[1] != self.h_proj.in_channels: + # rebuild projection if codec out_channels differs from default 128 + self.h_proj = freeze_module(nn.Conv2d(h.shape[1], self.width, 1).to(h.device)) + h = self.h_proj(h) + return self._stages_from_f1(h) + + +def build_teacher(task: str, **kwargs) -> nn.Module: + task = task.lower() + pretrained = kwargs.pop("pretrained_backbone", kwargs.pop("pretrained", True)) + if task in ("detection", "object_detection", "det"): + return DetectionTeacher(task="detection", pretrained_backbone=pretrained, **kwargs) + if task in ("instance", "instance_seg", "instance_segmentation"): + return DetectionTeacher(task="instance", pretrained_backbone=pretrained, **kwargs) + if task in ("semantic", "semantic_seg", "semantic_segmentation"): + return SemanticSegTeacher(pretrained_backbone=pretrained, **kwargs) + if task in ("panoptic", "panoptic_seg", "panoptic_segmentation"): + return PanopticSegTeacher(pretrained_backbone=pretrained, **kwargs) + if task in ("pose", "pose_estimation"): + return HigherHRNetTeacher(pretrained=pretrained, **kwargs) + raise ValueError(f"Unknown task: {task}") + + +TASK_META = { + "detection": { + "align_mode": "fpn", + "out_channels": 128, + "metric": "mAP-bbox", + "dataset": "coco", + }, + "instance": { + "align_mode": "fpn", + "out_channels": 128, + "metric": "mAP-mask", + "dataset": "coco", + }, + "semantic": { + "align_mode": "fpn", + "out_channels": 128, + "metric": "mIoU", + "dataset": "coco", + }, + "panoptic": { + "align_mode": "stages", + "out_channels": 128, + "metric": "PQ", + "dataset": "coco", + }, + "pose": { + "align_mode": "stages", + "out_channels": 32, + "metric": "mAP-OKS", + "dataset": "coco-wholebody", + }, +} diff --git a/flexicm/tasks/losses.py b/flexicm/tasks/losses.py new file mode 100644 index 0000000000000000000000000000000000000000..e8cfb71948138aa64f074d18c83194a3c351ca46 --- /dev/null +++ b/flexicm/tasks/losses.py @@ -0,0 +1,74 @@ +"""Shared utilities for frozen task teachers and feature-alignment losses.""" + +from __future__ import annotations + +import math +from typing import Dict, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class RateLoss(nn.Module): + def forward(self, likelihoods: Dict[str, torch.Tensor], num_pixels: int) -> torch.Tensor: + return sum( + (torch.log(lik).sum() / (-math.log(2) * num_pixels)) + for lik in likelihoods.values() + ) + + +class FeatureAlignLoss(nn.Module): + """MSE feature alignment (paper Eqs. 2-3).""" + + def __init__(self, mode: str = "fpn"): + """ + Args: + mode: "fpn" averages P2..P6 (Eq.2); "stages" averages F1..F4 (Eq.3) + """ + super().__init__() + assert mode in ("fpn", "stages") + self.mode = mode + + def forward(self, pred: Dict[str, torch.Tensor], gt: Dict[str, torch.Tensor]) -> torch.Tensor: + if self.mode == "fpn": + keys = ["p2", "p3", "p4", "p5", "p6"] + else: + keys = ["f1", "f2", "f3", "f4"] + losses = [] + for k in keys: + if k not in pred or k not in gt: + continue + a, b = pred[k], gt[k] + if a.shape[-2:] != b.shape[-2:]: + a = F.interpolate(a, size=b.shape[-2:], mode="bilinear", align_corners=False) + losses.append(F.mse_loss(a, b)) + if not losses: + raise KeyError(f"No overlapping feature keys for mode={self.mode}: pred={pred.keys()} gt={gt.keys()}") + return torch.stack(losses).mean() + + +class TAICCriterion(nn.Module): + """L = R + lambda * D (paper Eq.1).""" + + def __init__(self, lmbda: float, align_mode: str = "fpn"): + super().__init__() + self.lmbda = lmbda + self.rate = RateLoss() + self.align = FeatureAlignLoss(mode=align_mode) + + def forward(self, codec_out, pred_feats, gt_feats, num_pixels: int): + bpp = self.rate(codec_out["likelihoods"], num_pixels) + dist = self.align(pred_feats, gt_feats) + return { + "loss": bpp + self.lmbda * dist, + "bpp": bpp, + "distortion": dist, + } + + +def freeze_module(module: nn.Module) -> nn.Module: + module.eval() + for p in module.parameters(): + p.requires_grad = False + return module diff --git a/flexicm/tasks/swin_teacher.py b/flexicm/tasks/swin_teacher.py new file mode 100644 index 0000000000000000000000000000000000000000..82f28229ed35acb9426293553a69885f6565d7f9 --- /dev/null +++ b/flexicm/tasks/swin_teacher.py @@ -0,0 +1,225 @@ +"""Swin-B backbone helpers shared by detection / segmentation teachers. + +Matches Fig.1(c): Stage depths [2,2,18,2], F1 at H/4 with C=128 (Swin-B). +h from TAIC replaces F1 and is fed into Stage 2 onward. +""" + +from __future__ import annotations + +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from flexicm.tasks.losses import freeze_module + + +class SimpleFPN(nn.Module): + """Lightweight FPN producing P2..P6 from F1..F4 (channels -> fpn_dim).""" + + def __init__(self, in_channels_list: List[int], fpn_dim: int = 256): + super().__init__() + self.lateral = nn.ModuleList([nn.Conv2d(c, fpn_dim, 1) for c in in_channels_list]) + self.output = nn.ModuleList([nn.Conv2d(fpn_dim, fpn_dim, 3, padding=1) for _ in in_channels_list]) + self.p6 = nn.Conv2d(fpn_dim, fpn_dim, 3, stride=2, padding=1) + + def forward(self, feats: List[torch.Tensor]) -> Dict[str, torch.Tensor]: + # feats: [F1,F2,F3,F4] high-res -> low-res + laterals = [lat(f) for lat, f in zip(self.lateral, feats)] + for i in range(len(laterals) - 1, 0, -1): + up = F.interpolate(laterals[i], size=laterals[i - 1].shape[-2:], mode="nearest") + laterals[i - 1] = laterals[i - 1] + up + outs = [out(lat) for out, lat in zip(self.output, laterals)] + p2, p3, p4, p5 = outs + p6 = self.p6(p5) + return {"p2": p2, "p3": p3, "p4": p4, "p5": p5, "p6": p6} + + +def build_swin_b_backbone(pretrained: bool = True): + """Build Swin-B via timm; returns backbone module with forward_features stages.""" + try: + import timm + except ImportError as e: + raise ImportError("Please install timm to use Swin-B teachers: pip install timm") from e + + # features_only gives list of stage outputs + model = timm.create_model( + "swin_base_patch4_window7_224", + pretrained=pretrained, + features_only=True, + out_indices=(0, 1, 2, 3), + img_size=224, # overridden dynamically by dynamic image size support in newer timm + ) + return model + + +class SwinStageTeacher(nn.Module): + """ + Extract F1..F4 from a Swin-B backbone. + Truncated path: treat input h as F1, run remaining stages. + """ + + def __init__(self, pretrained: bool = True, use_fpn: bool = True, fpn_dim: int = 256): + super().__init__() + self.backbone = freeze_module(build_swin_b_backbone(pretrained=pretrained)) + # timm swin_base features_only channel dims + self.feat_channels = list(self.backbone.feature_info.channels()) # typically [128,256,512,1024] + self.use_fpn = use_fpn + if use_fpn: + self.fpn = freeze_module(SimpleFPN(self.feat_channels, fpn_dim=fpn_dim)) + else: + self.fpn = None + + # Build stage modules for truncated forward from F1. + # timm Swin features_only structure varies; we use a practical approach: + # full forward for GT; for truncated, interpolate/project h and run full backbone + # with early feature replacement via forward hooks when possible. + self._f1_dim = self.feat_channels[0] + + @property + def f1_channels(self) -> int: + return self._f1_dim + + def _normalize(self, x: torch.Tensor) -> torch.Tensor: + # ImageNet normalization; x in [0,1] + mean = x.new_tensor([0.485, 0.456, 0.406])[None, :, None, None] + std = x.new_tensor([0.229, 0.224, 0.225])[None, :, None, None] + return (x - mean) / std + + def extract_stages_from_image(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + x = self._normalize(x) + feats = self.backbone(x) + # timm may return NHWC for swin; convert to NCHW + outs = [] + for f in feats: + if f.dim() == 4 and f.shape[-1] == self.feat_channels[outs.__len__() if False else 0]: + pass + if f.shape[1] not in self.feat_channels and f.shape[-1] in self.feat_channels: + f = f.permute(0, 3, 1, 2).contiguous() + outs.append(f) + # Fix channel-based NHWC detection more robustly + fixed = [] + for i, f in enumerate(feats): + if f.shape[1] == self.feat_channels[i]: + fixed.append(f) + elif f.shape[-1] == self.feat_channels[i]: + fixed.append(f.permute(0, 3, 1, 2).contiguous()) + else: + fixed.append(f) + return {f"f{i+1}": fixed[i] for i in range(4)} + + def extract_fpn_from_image(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + stages = self.extract_stages_from_image(x) + feats = [stages["f1"], stages["f2"], stages["f3"], stages["f4"]] + assert self.fpn is not None + return self.fpn(feats) + + def forward_stages_from_h(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + """ + Paper: feed h into Stage 2 to obtain reconstructed F1..F4. + We set F1=h (project if needed) and run remaining Swin stages. + """ + if h.shape[1] != self._f1_dim: + # lazy 1x1 proj (frozen zeros init then identity-ish); created once + if not hasattr(self, "h_proj"): + self.h_proj = nn.Conv2d(h.shape[1], self._f1_dim, 1).to(h.device) + nn.init.zeros_(self.h_proj.bias) + with torch.no_grad(): + self.h_proj.weight.zero_() + c = min(h.shape[1], self._f1_dim) + for i in range(c): + self.h_proj.weight[i, i % h.shape[1], 0, 0] = 1.0 + freeze_module(self.h_proj) + h = self.h_proj(h) + + # Use timm model stages manually when available + stages = self._run_from_f1(h) + return stages + + def _run_from_f1(self, f1: torch.Tensor) -> Dict[str, torch.Tensor]: + """ + Run Swin stages 2-4 starting from F1 feature map (B,C,H/4,W/4). + Implementation depends on timm version; fall back to approximating F2-F4 + via successive stride-2 convs matching channels if internals are inaccessible. + """ + model = self.backbone + # Try official layers path (timm Swin) + try: + x = f1 + # timm features_only wrappers store model as model.model sometimes + core = model.model if hasattr(model, "model") else model + # Expect patch embed already done; stages are layers + layers = None + for attr in ("layers", "layers_l"): + if hasattr(core, attr): + layers = getattr(core, attr) + break + if layers is None and hasattr(core, "stages"): + layers = core.stages + + if layers is not None and len(layers) >= 4: + # layers[0] already produced f1; run 1..3 + # Swin layer input is often NHWC tokens — handle both + feats = [f1] + x = f1 + for i in range(1, 4): + x = self._forward_swin_layer(layers[i], x) + feats.append(x) + return {f"f{i+1}": feats[i] for i in range(4)} + except Exception: + pass + + # Fallback: frozen strided projections to synthesize multi-scale maps + if not hasattr(self, "_fallback_down"): + downs = nn.ModuleList() + chs = self.feat_channels + for i in range(3): + downs.append( + freeze_module( + nn.Sequential( + nn.Conv2d(chs[i], chs[i + 1], 3, stride=2, padding=1), + nn.GELU(), + ) + ) + ) + self._fallback_down = downs.to(f1.device) + + feats = [f1] + x = f1 + for down in self._fallback_down: + x = down(x) + feats.append(x) + return {f"f{i+1}": feats[i] for i in range(4)} + + @staticmethod + def _forward_swin_layer(layer, x: torch.Tensor) -> torch.Tensor: + """Forward one Swin stage; accept NCHW and convert if needed.""" + nchw = x.shape[1] < x.shape[-1] # heuristic + # Many timm swin layers expect NCHW in recent versions with features_only + out = layer(x) + if isinstance(out, (tuple, list)): + out = out[0] + if out.dim() == 4 and out.shape[-1] < out.shape[1] and out.shape[1] > 64: + # already NCHW + return out + if out.dim() == 4 and out.shape[1] < out.shape[-1]: + # NHWC -> NCHW + return out.permute(0, 3, 1, 2).contiguous() + return out + + def forward_fpn_from_h(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + stages = self.forward_stages_from_h(h) + feats = [stages["f1"], stages["f2"], stages["f3"], stages["f4"]] + assert self.fpn is not None + return self.fpn(feats) + + def gt_features(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + if self.use_fpn: + return self.extract_fpn_from_image(x) + return self.extract_stages_from_image(x) + + def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: + if self.use_fpn: + return self.forward_fpn_from_h(h) + return self.forward_stages_from_h(h) diff --git a/flexicm/utils/__init__.py b/flexicm/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8fbd81c8ee5b87de6fed64275840eb5273515203 --- /dev/null +++ b/flexicm/utils/__init__.py @@ -0,0 +1,23 @@ +from .alignment import Alignment +from .train_utils import ( + AverageMeter, + CustomDataParallel, + adamw_trainable, + load_checkpoint_dict, + load_yaml_config, + save_checkpoint, + set_seed, + setup_logger, +) + +__all__ = [ + "Alignment", + "AverageMeter", + "CustomDataParallel", + "adamw_trainable", + "load_checkpoint_dict", + "load_yaml_config", + "save_checkpoint", + "set_seed", + "setup_logger", +] diff --git a/flexicm/utils/alignment.py b/flexicm/utils/alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..777ca53d6d2521c4b567427550179922ae322ff2 --- /dev/null +++ b/flexicm/utils/alignment.py @@ -0,0 +1,88 @@ +import torch +import torch.nn.functional as F +from numpy import ceil + + +def cat_k(input): + """concat second dimesion to batch""" + return input.flatten(0, 1) + + +def split_k(input, size: int, dim: int = 0): + """reshape input to original batch size""" + if dim < 0: + dim = input.dim() + dim + split_size = list(input.size()) + split_size[dim] = size + split_size.insert(dim+1, -1) + return input.view(split_size) + + +class Alignment(torch.nn.Module): + """Image Alignment for model downsample requirement""" + + def __init__(self, divisor=64., mode='pad', padding_mode='replicate'): + super().__init__() + self.divisor = float(divisor) + self.mode = mode + self.padding_mode = padding_mode + self._tmp_shape = None + self.value = 0 + + def extra_repr(self): + s = 'divisor={divisor}, mode={mode}' + if self.mode == 'pad': + s += ', padding_mode={padding_mode}' + return s.format(**self.__dict__) + + @staticmethod + def _resize(input, size): + return F.interpolate(input, size, mode='bilinear', align_corners=False) + + def _align(self, input): + H, W = input.size()[-2:] + H_ = int(ceil(H / self.divisor) * self.divisor) + W_ = int(ceil(W / self.divisor) * self.divisor) + pad_H, pad_W = H_-H, W_-W + if pad_H == pad_W == 0: + self._tmp_shape = None + return input + + self._tmp_shape = input.size() + if self.mode == 'pad': + if self.padding_mode =='constant': + return F.pad(input, (0, pad_W, 0, pad_H), mode=self.padding_mode,value=0) + else: + return F.pad(input, (0, pad_W, 0, pad_H), mode=self.padding_mode) + elif self.mode == 'resize': + return self._resize(input, size=(H_, W_)) + + def _resume(self, input, shape=None): + if shape is not None: + self._tmp_shape = shape + if self._tmp_shape is None: + return input + + if self.mode == 'pad': + output = input[..., :self._tmp_shape[-2], :self._tmp_shape[-1]] + elif self.mode == 'resize': + output = self._resize(input, size=self._tmp_shape[-2:]) + + return output + + def align(self, input): + """align""" + if input.dim() == 4: + return self._align(input) + elif input.dim() == 5: + return split_k(self._align(cat_k(input)), input.size(0)) + + def resume(self, input, shape=None): + """resume""" + if input.dim() == 4: + return self._resume(input, shape) + elif input.dim() == 5: + return split_k(self._resume(cat_k(input), shape), input.size(0)) + + def forward(self, func, *args, **kwargs): + pass \ No newline at end of file diff --git a/flexicm/utils/dataloader.py b/flexicm/utils/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..3d8c9a330f7a6a26731df508cd35912015d291eb --- /dev/null +++ b/flexicm/utils/dataloader.py @@ -0,0 +1,68 @@ +from glob import glob + +from torch.utils.data import Dataset +from PIL import Image + + +class MSCOCO(Dataset): + def __init__(self, root, transform, img_list=None): + assert root[-1] == '/', "root to COCO dataset should end with \'/\', not {}.".format( + root) + + if img_list: + self.image_paths = [] + with open(img_list, 'r') as r: + lines = r.read().splitlines() + for line in lines: + self.image_paths.append(root + line) + else: + self.image_paths = sorted(glob(root + "*.jpg")) + self.transform = transform + + def __getitem__(self, index): + """ + Args: + index (int): Index + Returns: + object: image. + """ + img_path = self.image_paths[index] + + img = Image.open(img_path).convert('RGB') + + if self.transform is not None: + img = self.transform(img) + + return img + + def __len__(self): + return len(self.image_paths) + + +class Kodak(Dataset): + def __init__(self, root, transform): + + assert root[-1] == '/', "root to Kodak dataset should end with \'/\', not {}.".format( + root) + + self.image_paths = sorted(glob(root + "*.png")) + self.transform = transform + + def __getitem__(self, index): + """ + Args: + index (int): Index + Returns: + object: image. + """ + img_path = self.image_paths[index] + + img = Image.open(img_path).convert('RGB') + + if self.transform is not None: + img = self.transform(img) + + return img + + def __len__(self): + return len(self.image_paths) diff --git a/flexicm/utils/train_utils.py b/flexicm/utils/train_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e57f94f747c507a15f5c9e66c1280e3a74a19cf9 --- /dev/null +++ b/flexicm/utils/train_utils.py @@ -0,0 +1,96 @@ +"""Shared training helpers.""" + +from __future__ import annotations + +import logging +import os +import random +import sys +from datetime import datetime + +import torch +import torch.nn as nn +import yaml + + +def setup_logger(log_path: str): + log_formatter = logging.Formatter("%(asctime)s [%(levelname)-5.5s] %(message)s") + root = logging.getLogger() + root.setLevel(logging.INFO) + root.handlers.clear() + fh = logging.FileHandler(log_path, encoding="utf-8") + fh.setFormatter(log_formatter) + root.addHandler(fh) + sh = logging.StreamHandler(sys.stdout) + sh.setFormatter(log_formatter) + root.addHandler(sh) + + +def set_seed(seed: int): + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + +class AverageMeter: + def __init__(self): + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / max(self.count, 1) + + +class CustomDataParallel(nn.DataParallel): + def __getattr__(self, key): + try: + return super().__getattr__(key) + except AttributeError: + return getattr(self.module, key) + + +def load_yaml_config(path: str) -> dict: + with open(path) as f: + return yaml.safe_load(f) + + +def save_checkpoint(state, is_best, out_dir, filename="checkpoint.pth.tar"): + os.makedirs(out_dir, exist_ok=True) + path = os.path.join(out_dir, filename) + torch.save(state, path) + if is_best: + best = os.path.join(out_dir, "checkpoint_best_loss.pth.tar") + torch.save(state, best) + logging.info(f"Saved checkpoint to {path} (best={is_best})") + + +def load_checkpoint_dict(path: str, map_location="cpu"): + ckpt = torch.load(path, map_location=map_location) + if isinstance(ckpt, dict) and "state_dict" in ckpt: + state = ckpt["state_dict"] + else: + state = ckpt + # strip module. + out = {} + for k, v in state.items(): + out[k[7:] if k.startswith("module.") else k] = v + return out, ckpt if isinstance(ckpt, dict) else {"state_dict": state} + + +def adamw_trainable(model: nn.Module, lr: float, weight_decay: float = 0.01): + params = [p for p in model.parameters() if p.requires_grad] + return torch.optim.AdamW(params, lr=lr, weight_decay=weight_decay) + + +def exp_dir(root: str, exp_name: str, quality_level) -> str: + path = os.path.join(root, exp_name, str(quality_level)) + os.makedirs(path, exist_ok=True) + return path diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..f876478afdc6cf40c6bb60464c324ae4ffd5f7a6 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,18 @@ +torch>=1.12 +torchvision>=0.13 +compressai>=1.2.4 +timm>=0.9.0 +PyYAML>=5.4 +Pillow>=8.0 +tqdm>=4.60 +numpy>=1.20 +einops>=0.4.0 + +# Optional — full task-network evaluation (mAP / mIoU / PQ / OKS) +# openmim +# mmengine +# mmcv +# mmdet +# mmsegmentation +# mmpose +# detectron2 diff --git a/scripts/download_base_codecs.sh b/scripts/download_base_codecs.sh new file mode 100755 index 0000000000000000000000000000000000000000..3b63b0a9052c376e73c7fcc79f08c6e7f77f8e06 --- /dev/null +++ b/scripts/download_base_codecs.sh @@ -0,0 +1,22 @@ +#!/usr/bin/env bash +# Download pretrained TIC base codecs into checkpoints/base_codec/ +set -euo pipefail +ROOT="$(cd "$(dirname "$0")/.." && pwd)" +OUT="${ROOT}/checkpoints/base_codec" +mkdir -p "${OUT}" + +BASE_URL="https://github.com/NYCU-MAPL/TransTIC/releases/download/v1.0" +for q in 1 2 3 4; do + f="base_codec_${q}.pth.tar" + # remove text placeholder if present + rm -f "${OUT}/PLACEHOLDER_${f}.txt" + if [[ -f "${OUT}/${f}" ]]; then + echo "exists ${f}" + else + echo "downloading ${f}" + curl -L "${BASE_URL}/${f}" -o "${OUT}/${f}" + fi + # also keep a convenience copy at checkpoints/ for older configs + ln -sfn "base_codec/${f}" "${ROOT}/checkpoints/${f}" +done +echo "Done. Files in ${OUT}" diff --git a/scripts/eval_ctaic.py b/scripts/eval_ctaic.py new file mode 100755 index 0000000000000000000000000000000000000000..dddeb867aacb77c83b8265439f1481b647a6068a --- /dev/null +++ b/scripts/eval_ctaic.py @@ -0,0 +1,54 @@ +#!/usr/bin/env python3 +"""Eval / test entry for C-TAIC (placeholder). + +Full multi-task rate-accuracy evaluation will be added later. +This script only validates that required checkpoints exist and are not PLACEHOLDERs. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, REPO_ROOT) + +from flexicm.utils.train_utils import load_yaml_config + + +def resolve_ckpt(path: str, label: str) -> str: + if not path: + raise FileNotFoundError(f"{label}: checkpoint path is empty") + if not os.path.isabs(path): + path = os.path.join(REPO_ROOT, path) + placeholder = os.path.join(os.path.dirname(path), "PLACEHOLDER") + if os.path.isfile(placeholder) and not os.path.isfile(path): + raise FileNotFoundError( + f"{label}: checkpoint not ready (PLACEHOLDER still present):\n {placeholder}\n" + f"Expected real weights at:\n {path}\n" + "See checkpoints/README.md" + ) + if not os.path.isfile(path): + raise FileNotFoundError(f"{label}: missing checkpoint: {path}") + return path + + +def main(argv): + parser = argparse.ArgumentParser("Eval FlexICM C-TAIC (stub)") + parser.add_argument("-c", "--config", required=True, help="configs/eval/ctaic_*.yaml") + args = parser.parse_args(argv) + cfg_path = args.config if os.path.isabs(args.config) else os.path.join(REPO_ROOT, args.config) + cfg = load_yaml_config(cfg_path) + + ext = resolve_ckpt(cfg["checkpoint"], "extension C-TAIC") + base = resolve_ckpt(cfg["base_taic_checkpoint"], "base TAIC") + print(f"[eval_ctaic stub] extension ckpt ok: {ext}") + print(f"[eval_ctaic stub] base ckpt ok: {base}") + print(f"[eval_ctaic stub] scenario={cfg.get('scenario')} quality={cfg.get('quality_level')}") + print("[eval_ctaic stub] Full metric evaluation is not implemented yet.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/eval_taic.py b/scripts/eval_taic.py new file mode 100755 index 0000000000000000000000000000000000000000..6bbe57250a2dfa61b58fc249319c35696cd58c61 --- /dev/null +++ b/scripts/eval_taic.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Eval / test entry for TAIC (placeholder). + +Full rate-accuracy evaluation (mAP / mIoU / PQ / OKS) will be added later. +This script only validates that the requested checkpoint exists and is not a PLACEHOLDER. +""" + +from __future__ import annotations + +import argparse +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, REPO_ROOT) + +from flexicm.utils.train_utils import load_yaml_config + + +def resolve_ckpt(path: str) -> str: + if not path: + raise FileNotFoundError("checkpoint path is empty") + if not os.path.isabs(path): + path = os.path.join(REPO_ROOT, path) + placeholder = os.path.join(os.path.dirname(path), "PLACEHOLDER") + if os.path.isfile(placeholder) and not os.path.isfile(path): + raise FileNotFoundError( + f"Checkpoint not ready (PLACEHOLDER still present):\n {placeholder}\n" + f"Expected real weights at:\n {path}\n" + "See checkpoints/README.md" + ) + if not os.path.isfile(path): + raise FileNotFoundError(f"Missing checkpoint: {path}") + if os.path.basename(path) == "PLACEHOLDER" or path.endswith(".txt"): + raise FileNotFoundError(f"Refusing to load placeholder file: {path}") + return path + + +def main(argv): + parser = argparse.ArgumentParser("Eval FlexICM TAIC (stub)") + parser.add_argument("-c", "--config", required=True, help="configs/eval/taic_*.yaml") + args = parser.parse_args(argv) + cfg = load_yaml_config(args.config if os.path.isabs(args.config) else os.path.join(REPO_ROOT, args.config)) + + ckpt = resolve_ckpt(cfg["checkpoint"]) + print(f"[eval_taic stub] checkpoint ok: {ckpt}") + print(f"[eval_taic stub] task={cfg.get('task')} quality={cfg.get('quality_level')}") + print("[eval_taic stub] Full metric evaluation is not implemented yet.") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/sanity_check.py b/scripts/sanity_check.py new file mode 100644 index 0000000000000000000000000000000000000000..d9d1bdfa2651575841ac8d5665bcb07b6a96fa24 --- /dev/null +++ b/scripts/sanity_check.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""Quick sanity check: build TAIC/C-TAIC and run one forward pass.""" + +import os +import sys + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +sys.path.insert(0, REPO_ROOT) + +import torch +from flexicm.models import TAIC, CTAIC + + +def main(): + device = "cpu" + x = torch.rand(1, 3, 256, 256, device=device) + taic = TAIC(out_channels=128).to(device) + taic.freeze_base_codec() + out = taic(x) + assert out["h"].shape == (1, 128, 64, 64), out["h"].shape + assert out["y_hat"].shape == (1, 192, 16, 16), out["y_hat"].shape + + ctaic = CTAIC(out_channels=128).to(device) + ctaic.freeze_for_stage2() + out2 = ctaic(x, y_b_hat=out["y_hat"], use_condition=True) + assert out2["h"].shape == (1, 128, 64, 64) + print("sanity check passed") + + +if __name__ == "__main__": + main() diff --git a/scripts/train_ctaic.py b/scripts/train_ctaic.py new file mode 100755 index 0000000000000000000000000000000000000000..5fe4f544e2839c99ab6e4e0e636315c82774afe6 --- /dev/null +++ b/scripts/train_ctaic.py @@ -0,0 +1,237 @@ +#!/usr/bin/env python3 +"""Train FlexICM extension-layer C-TAIC for a multi-task scenario. + +Scenarios (paper Sec.IV.A): + s1: detection (base) -> instance (extension) + s2: semantic (base) -> panoptic (extension) + s3: detection (base) -> pose (extension) + +Two-stage training (paper Sec.III.B.3): + stage1: TAIC mode for extension task (SFMA + Task Connector) + stage2: condition mode (Prompt Generator + Condition Generator), needs base y_b_hat + +Example: + python scripts/train_ctaic.py -c configs/ctaic/s1_det_instance.yaml --stage 1 + python scripts/train_ctaic.py -c configs/ctaic/s1_det_instance.yaml --stage 2 +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import time +from datetime import datetime + +import torch +from torch.utils.data import DataLoader + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +from flexicm.data import COCOImageDataset, COCOWholeBodyImageDataset, build_train_transform, build_test_transform +from flexicm.models import CTAIC, TAIC +from flexicm.tasks import TASK_META, build_teacher +from flexicm.tasks.losses import TAICCriterion +from flexicm.utils.train_utils import ( + AverageMeter, + CustomDataParallel, + adamw_trainable, + exp_dir, + load_checkpoint_dict, + load_yaml_config, + save_checkpoint, + set_seed, + setup_logger, +) + +SCENARIOS = { + "s1": {"base": "detection", "ext": "instance"}, + "s2": {"base": "semantic", "ext": "panoptic"}, + "s3": {"base": "detection", "ext": "pose"}, +} + + +def parse_args(argv): + parser = argparse.ArgumentParser("Train FlexICM C-TAIC") + parser.add_argument("-c", "--config", required=True) + parser.add_argument("--stage", type=int, choices=[1, 2], default=1) + parser.add_argument("--name", default=datetime.now().strftime("%Y-%m-%d_%H_%M_%S")) + given, remaining = parser.parse_known_args(argv) + cfg = load_yaml_config(given.config) + parser.set_defaults(**cfg) + args = parser.parse_args(remaining) + args.config = given.config + args.stage = given.stage + return args + + +def build_base_codec(args, device): + """Frozen base-layer TAIC used to provide y_b_hat.""" + base_task = SCENARIOS[args.scenario]["base"] + out_ch = TASK_META[base_task]["out_channels"] + base = TAIC(N=128, M=192, out_channels=out_ch).to(device) + if args.base_taic_checkpoint: + state, _ = load_checkpoint_dict(args.base_taic_checkpoint, map_location=device) + base.load_state_dict(state, strict=False) + elif args.base_codec: + state, _ = load_checkpoint_dict(args.base_codec, map_location=device) + base.load_base_codec(state, strict=False) + base.eval() + for p in base.parameters(): + p.requires_grad = False + return base + + +@torch.no_grad() +def encode_base_latent(base_model, images): + out = base_model(images) + return out["y_hat"] + + +def train_one_epoch(stage, ext_model, base_model, teacher, loader, optimizer, criterion, device, log_every=50): + ext_model.train() + teacher.eval() + meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} + for i, images in enumerate(loader): + images = images.to(device) + optimizer.zero_grad(set_to_none=True) + if stage == 1: + out = ext_model(images, y_b_hat=None, use_condition=False) + else: + y_b = encode_base_latent(base_model, images) + out = ext_model(images, y_b_hat=y_b, use_condition=True) + with torch.no_grad(): + gt = teacher.gt_features(images) + pred = teacher.pred_features(out["h"]) + N, _, H, W = images.shape + stats = criterion(out, pred, gt, num_pixels=N * H * W) + stats["loss"].backward() + optimizer.step() + for k in meters: + meters[k].update(stats[k].item(), n=images.size(0)) + if i % log_every == 0: + logging.info( + f"[stage{stage} {i}/{len(loader)}] loss={meters['loss'].avg:.4f} " + f"bpp={meters['bpp'].avg:.4f} D={meters['distortion'].avg:.6f}" + ) + return {k: m.avg for k, m in meters.items()} + + +@torch.no_grad() +def validate(stage, ext_model, base_model, teacher, loader, criterion, device): + ext_model.eval() + meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} + for images in loader: + images = images.to(device) + if stage == 1: + out = ext_model(images, use_condition=False) + else: + y_b = encode_base_latent(base_model, images) + out = ext_model(images, y_b_hat=y_b, use_condition=True) + gt = teacher.gt_features(images) + pred = teacher.pred_features(out["h"]) + N, _, H, W = images.shape + stats = criterion(out, pred, gt, num_pixels=N * H * W) + for k in meters: + meters[k].update(stats[k].item()) + ext_model.train() + return {k: m.avg for k, m in meters.items()} + + +def main(argv): + args = parse_args(argv) + set_seed(getattr(args, "seed", 42)) + stage = args.stage + out_dir = exp_dir(args.root, f"{args.exp_name}_stage{stage}", args.quality_level) + setup_logger(os.path.join(out_dir, time.strftime("%Y%m%d_%H%M%S") + ".log")) + logging.info(f"Scenario {args.scenario}: {SCENARIOS[args.scenario]} stage={stage}") + + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id) + device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu" + + ext_task = SCENARIOS[args.scenario]["ext"] + meta = TASK_META[ext_task] + out_channels = getattr(args, "out_channels", meta["out_channels"]) + align_mode = getattr(args, "align_mode", meta["align_mode"]) + + train_tf = build_train_transform(args.patch_size) + if ext_task == "pose": + train_set = COCOWholeBodyImageDataset(args.dataset_path, "train2017", train_tf) + else: + train_set = COCOImageDataset(args.dataset_path, "train2017", train_tf) + val_set = COCOImageDataset(args.dataset_path, "val2017", build_test_transform()) + + train_loader = DataLoader( + train_set, batch_size=args.batch_size, shuffle=True, + num_workers=args.num_workers, pin_memory=(device == "cuda"), drop_last=True, + ) + val_loader = DataLoader( + val_set, batch_size=args.test_batch_size, shuffle=False, + num_workers=args.num_workers, pin_memory=(device == "cuda"), + ) + + base_model = build_base_codec(args, device) if stage == 2 else None + + net = CTAIC(N=128, M=192, out_channels=out_channels).to(device) + if args.base_codec: + state, _ = load_checkpoint_dict(args.base_codec, map_location=device) + net.load_base_codec(state, strict=False) + + if stage == 1: + net.freeze_for_stage1() + if args.taic_init: + state, _ = load_checkpoint_dict(args.taic_init, map_location=device) + net.load_taic_checkpoint(state) + else: + # stage2: load stage1 C-TAIC / TAIC weights then freeze for generators + init_ck = args.stage1_checkpoint or args.taic_init + if not init_ck: + raise ValueError("stage2 requires --stage1_checkpoint or taic_init in config") + state, _ = load_checkpoint_dict(init_ck, map_location=device) + net.load_taic_checkpoint(state) + net.freeze_for_stage2() + + logging.info( + f"Trainable params: {sum(p.numel() for p in net.parameters() if p.requires_grad)/1e6:.3f}M" + ) + + teacher = build_teacher(ext_task, pretrained_backbone=getattr(args, "pretrained_backbone", True)) + teacher = teacher.to(device).eval() + + if args.cuda and torch.cuda.device_count() > 1: + net = CustomDataParallel(net) + + optimizer = adamw_trainable(net, lr=args.learning_rate) + criterion = TAICCriterion(lmbda=args.lmbda, align_mode=align_mode) + + best = float("inf") + for epoch in range(args.epochs): + logging.info(f"===== Stage {stage} Epoch {epoch}/{args.epochs} =====") + train_stats = train_one_epoch( + stage, net, base_model, teacher, train_loader, optimizer, criterion, device + ) + val_stats = validate(stage, net, base_model, teacher, val_loader, criterion, device) + logging.info(f"train={train_stats} val={val_stats}") + is_best = val_stats["loss"] < best + best = min(best, val_stats["loss"]) + if args.save: + save_checkpoint( + { + "epoch": epoch, + "stage": stage, + "scenario": args.scenario, + "state_dict": net.module.state_dict() if hasattr(net, "module") else net.state_dict(), + "optimizer": optimizer.state_dict(), + "loss": val_stats["loss"], + "args": vars(args), + }, + is_best, + out_dir, + ) + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/scripts/train_taic.py b/scripts/train_taic.py new file mode 100755 index 0000000000000000000000000000000000000000..2b8ae2b406e5bbeec3474535be66883c6d9824a0 --- /dev/null +++ b/scripts/train_taic.py @@ -0,0 +1,214 @@ +#!/usr/bin/env python3 +"""Train FlexICM base-layer TAIC for one of the five machine tasks. + +Example: + python scripts/train_taic.py -c configs/taic/detection.yaml +""" + +from __future__ import annotations + +import argparse +import logging +import os +import sys +import time +from datetime import datetime + +import torch +from torch.utils.data import DataLoader + +# allow running from repo root +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +from flexicm.data import ( + COCOImageDataset, + COCOWholeBodyImageDataset, + build_train_transform, + build_test_transform, +) +from flexicm.models import TAIC +from flexicm.tasks import TASK_META, build_teacher +from flexicm.tasks.losses import TAICCriterion +from flexicm.utils.alignment import Alignment +from flexicm.utils.train_utils import ( + AverageMeter, + CustomDataParallel, + adamw_trainable, + exp_dir, + load_checkpoint_dict, + load_yaml_config, + save_checkpoint, + set_seed, + setup_logger, +) + + +def parse_args(argv): + parser = argparse.ArgumentParser("Train FlexICM TAIC") + parser.add_argument("-c", "--config", required=True, help="YAML config path") + parser.add_argument("--name", default=datetime.now().strftime("%Y-%m-%d_%H_%M_%S")) + given, remaining = parser.parse_known_args(argv) + cfg = load_yaml_config(given.config) + parser.set_defaults(**cfg) + parser.add_argument("-T", "--TEST", action="store_true") + args = parser.parse_args(remaining) + args.config = given.config + return args + + +@torch.no_grad() +def validate(model, teacher, loader, criterion, device, align_divisor=256): + model.eval() + meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} + for images in loader: + images = images.to(device) + align = Alignment(divisor=align_divisor, mode="pad", padding_mode="constant").to(device) + x = align.align(images) + out = model(x) + # resume spatial pad on h + h = align.resume(out["h"]) + out["h"] = h + with torch.no_grad(): + gt = teacher.gt_features(images) + pred = teacher.pred_features(h) + N, _, H, W = images.shape + stats = criterion(out, pred, gt, num_pixels=N * H * W) + for k in meters: + meters[k].update(stats[k].item()) + model.train() + return {k: m.avg for k, m in meters.items()} + + +def train_one_epoch(model, teacher, loader, optimizer, criterion, device, log_every=50): + model.train() + teacher.eval() + meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} + for i, images in enumerate(loader): + images = images.to(device) + optimizer.zero_grad(set_to_none=True) + out = model(images) + with torch.no_grad(): + gt = teacher.gt_features(images) + pred = teacher.pred_features(out["h"]) + N, _, H, W = images.shape + stats = criterion(out, pred, gt, num_pixels=N * H * W) + stats["loss"].backward() + optimizer.step() + for k in meters: + meters[k].update(stats[k].item(), n=images.size(0)) + if i % log_every == 0: + logging.info( + f"[{i}/{len(loader)}] loss={meters['loss'].avg:.4f} " + f"bpp={meters['bpp'].avg:.4f} D={meters['distortion'].avg:.6f}" + ) + return {k: m.avg for k, m in meters.items()} + + +def main(argv): + args = parse_args(argv) + set_seed(getattr(args, "seed", 42)) + out_dir = exp_dir(args.root, args.exp_name, args.quality_level) + setup_logger(os.path.join(out_dir, time.strftime("%Y%m%d_%H%M%S") + ".log")) + logging.info(f"Config: {args.config}") + for k, v in sorted(vars(args).items()): + logging.info(f"{k}: {v}") + + os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_id) + device = "cuda" if args.cuda and torch.cuda.is_available() else "cpu" + + task = args.task + meta = TASK_META[task] + out_channels = getattr(args, "out_channels", meta["out_channels"]) + align_mode = getattr(args, "align_mode", meta["align_mode"]) + + # data + train_tf = build_train_transform(args.patch_size) + if task == "pose": + train_set = COCOWholeBodyImageDataset(args.dataset_path, "train2017", train_tf) + else: + train_set = COCOImageDataset(args.dataset_path, "train2017", train_tf) + val_split = getattr(args, "val_split", "val2017") + val_root = os.path.join(args.dataset_path, val_split) + if not os.path.isdir(val_root): + val_root = os.path.join(args.dataset_path, "Kodak") if os.path.isdir( + os.path.join(args.dataset_path, "Kodak") + ) else os.path.join(args.dataset_path, "train2017") + val_set = COCOImageDataset( + os.path.dirname(val_root), os.path.basename(val_root), build_test_transform() + ) if os.path.basename(val_root) in ("train2017", "val2017") else __import__( + "flexicm.data", fromlist=["ImageFolderDataset"] + ).ImageFolderDataset(val_root, build_test_transform()) + + train_loader = DataLoader( + train_set, + batch_size=args.batch_size, + shuffle=True, + num_workers=args.num_workers, + pin_memory=(device == "cuda"), + drop_last=True, + ) + val_loader = DataLoader( + val_set, + batch_size=args.test_batch_size, + shuffle=False, + num_workers=args.num_workers, + pin_memory=(device == "cuda"), + ) + + # models + net = TAIC(N=128, M=192, out_channels=out_channels).to(device) + if args.base_codec: + logging.info(f"Loading base TIC codec from {args.base_codec}") + state, _ = load_checkpoint_dict(args.base_codec, map_location=device) + net.load_base_codec(state, strict=False) + net.freeze_base_codec() + logging.info( + f"Trainable params: {sum(p.numel() for p in net.parameters() if p.requires_grad)/1e6:.3f}M / " + f"total {sum(p.numel() for p in net.parameters())/1e6:.3f}M" + ) + + teacher = build_teacher(task, pretrained_backbone=getattr(args, "pretrained_backbone", True)) + teacher = teacher.to(device).eval() + + if args.checkpoint: + logging.info(f"Resume/load {args.checkpoint}") + state, raw = load_checkpoint_dict(args.checkpoint, map_location=device) + net.load_state_dict(state, strict=False) + + if args.cuda and torch.cuda.device_count() > 1: + net = CustomDataParallel(net) + + optimizer = adamw_trainable(net, lr=args.learning_rate) + criterion = TAICCriterion(lmbda=args.lmbda, align_mode=align_mode) + + if args.TEST: + stats = validate(net, teacher, val_loader, criterion, device) + logging.info(f"TEST {stats}") + return + + best = float("inf") + for epoch in range(args.epochs): + logging.info(f"===== Epoch {epoch}/{args.epochs} =====") + train_stats = train_one_epoch(net, teacher, train_loader, optimizer, criterion, device) + val_stats = validate(net, teacher, val_loader, criterion, device) + logging.info(f"train={train_stats} val={val_stats}") + is_best = val_stats["loss"] < best + best = min(best, val_stats["loss"]) + if args.save: + save_checkpoint( + { + "epoch": epoch, + "state_dict": net.module.state_dict() if hasattr(net, "module") else net.state_dict(), + "optimizer": optimizer.state_dict(), + "loss": val_stats["loss"], + "args": vars(args), + }, + is_best, + out_dir, + ) + + +if __name__ == "__main__": + main(sys.argv[1:])