diff --git a/.gitignore b/.gitignore deleted file mode 100644 index 2e4cb4492a0c066dd38bb99afccfd4c94421cee0..0000000000000000000000000000000000000000 --- a/.gitignore +++ /dev/null @@ -1,32 +0,0 @@ -__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 - -# Heavy weights to keep local-only. -# TAIC / C-TAIC under checkpoints/{taic,ctaic} are intentionally uploadable -# (avoid broad checkpoints/** + !negation — Hub upload handles ! poorly). -checkpoints/base_codec/**/*.pth.tar -checkpoints/base_codec/**/*.tar -checkpoints/base_codec/**/*.pkl -checkpoints/base_codec/**/*.pth -checkpoints/task_networks/**/*.pth.tar -checkpoints/task_networks/**/*.tar -checkpoints/task_networks/**/*.pkl -checkpoints/task_networks/**/*.pth -!checkpoints/**/PLACEHOLDER* -!checkpoints/**/*.txt -!checkpoints/README.md diff --git a/README.md b/README.md index d28ec6d39543cd284785532bd1a319614cf7233d..1210cd8ae5e4bc1e912161785ca2cae3dc812004 100644 --- a/README.md +++ b/README.md @@ -8,439 +8,5 @@ tags: # 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 | Cascade 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 | Swin-B (`timm`) for FPN alignment; metric head = Cascade Mask R-CNN + Swin-B | `timm`; ImageNet Swin-B on first run | -| Semantic / Panoptic | Swin-B (`timm`) | `timm`; ImageNet Swin-B on first run | -| 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 weights -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 -``` - -Official detection / instance weights from -[Swin-Transformer-Object-Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection) -(see `configs/task_networks/README.md`): - -- **Cascade Mask R-CNN + Swin-B** (detection mAP-bbox **and** instance mAP-mask; same weights) -- **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 -│ ├── eval_taic.py # codec test (bpp / feature D) -│ └── eval_ctaic.py # codec test for C-TAIC -├── 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 / codec-test configs: `configs/eval/`. - ---- - - - -## 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 -``` - ---- - - - -## Codec Test - -Codec test measures **compression statistics**: - - -| Metric | Meaning | -| ------------ | ----------------------------------------------------------------- | -| `bpp` | Likelihood bitrate R | -| `distortion` | Feature alignment D (Eq. 2 or Eq. 3) | -| `loss` | R + \lambda D | -| `actual_bpp` | Optional: real bitstream size after `compress()` / `decompress()` | - - -For C-TAIC, reported `bpp` is **extension-layer only** (base-layer rate is excluded), matching the paper. - -### Prepare codec checkpoints - -1. Copy trained weights into `checkpoints/taic/` or `checkpoints/ctaic/` (see `checkpoints/README.md`) -2. Remove the local `PLACEHOLDER` once `checkpoint_best_loss.pth.tar` is present -3. Edit `dataset_path` / `gpu_id` in `configs/eval/*.yaml` - -```bash -python scripts/eval_taic.py -c configs/eval/taic_detection.yaml -python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml -``` - ---- - - - -## Task-network metric evaluation - -To reproduce paper rate–accuracy numbers you must **also** load the official pretrained -**task networks** and run metrics on COCO val: - - -| Task | Task network | Metric | -| --------- | ---------------------------- | -------- | -| Detection | Cascade Mask R-CNN + Swin-B | mAP-bbox | -| Instance | Cascade Mask R-CNN + Swin-B | mAP-mask | -| Semantic | UPerNet + Swin-B | mIoU | -| Panoptic | MaskFormer + Swin-B | PQ | -| Pose | HigherHRNet (HRNet backbone) | mAP-OKS | - - -Pipeline: `image → codec → h → truncated task net (from Stage2 / FPN) → metric`. - -### Install metric dependencies - -```bash -pip install pycocotools -pip install -U openmim -mim install mmengine mmcv mmdet mmsegmentation mmpose -# optional for PQ: -# pip install git+https://github.com/cocodataset/panopticapi.git -``` - - - -### Prepare task-network configs & checkpoints - -1. Put / symlink real OpenMMLab configs under `configs/task_networks/` - (see `configs/task_networks/README.md`; current `*.py` files are stubs) -2. Download official weights to: - -```text -checkpoints/task_networks/ -├── detection/model.pth -├── instance/model.pth -├── semantic/model.pth -├── panoptic/model.pth -└── pose/model.pth -``` - -1. Set in each `configs/eval/*.yaml`: - -```yaml -task_config: "./configs/task_networks/.py" -task_checkpoint: "./checkpoints/task_networks//model.pth" -ann_file: "annotations/instances_val2017.json" -``` - - - -### Run codec + metrics - -```bash -python scripts/eval_taic.py -c configs/eval/taic_detection.yaml -python scripts/eval_taic.py -c configs/eval/taic_instance.yaml -python scripts/eval_taic.py -c configs/eval/taic_semantic.yaml -python scripts/eval_taic.py -c configs/eval/taic_panoptic.yaml -python scripts/eval_taic.py -c configs/eval/taic_pose.yaml - -python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml -python scripts/eval_ctaic.py -c configs/eval/ctaic_s2.yaml -python scripts/eval_ctaic.py -c configs/eval/ctaic_s3.yaml -``` - -JSON results (codec + task metrics) are written under `logs/eval_taic/` or `logs/eval_ctaic/`. - -> Detection / instance metric paths are the most complete (COCO bbox via pycocotools). -> Semantic mIoU needs a GT label loader; panoptic PQ needs `panopticapi` + GT folders; -> pose-from-`h` may need a HigherHRNet stem hook for your exact MMPose version. +Official checkpoints for the paper **FlexICM: A Flexible Image Coding for Machines Framework** (Tianma Shen, Ying Liu). diff --git a/checkpoints/README.md b/checkpoints/README.md deleted file mode 100644 index fe9792b459dd1944d00da5233989f23feeeef75e..0000000000000000000000000000000000000000 --- a/checkpoints/README.md +++ /dev/null @@ -1,92 +0,0 @@ -# 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. - -Codec test (bpp / feature distortion): - -```bash -python scripts/eval_taic.py -c configs/eval/taic_detection.yaml -python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml -``` - -Scripts refuse to run if a `PLACEHOLDER` file is still present or if the `.pth.tar` is missing. -Task rate–accuracy metrics are not implemented in these scripts yet. - -## Task networks (for metric evaluation) - -```text -checkpoints/task_networks/ -├── detection/model.pth -├── instance/model.pth -├── semantic/model.pth -├── panoptic/model.pth -└── pose/model.pth -``` - -These are **official pretrained task networks** (not codec weights). -Required when running: - -```bash -python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --with-metrics -``` - -See `configs/task_networks/README.md` for config/checkpoint pairing. diff --git a/checkpoints/base_codec/NOTE.txt b/checkpoints/base_codec/NOTE.txt deleted file mode 100644 index 3e35f8267fe1e54844bbb481277780a142b33648..0000000000000000000000000000000000000000 --- a/checkpoints/base_codec/NOTE.txt +++ /dev/null @@ -1,2 +0,0 @@ -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 deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/base_codec/PLACEHOLDER_base_codec_1.pth.tar.txt +++ /dev/null @@ -1,2 +0,0 @@ -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 deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/base_codec/PLACEHOLDER_base_codec_2.pth.tar.txt +++ /dev/null @@ -1,2 +0,0 @@ -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 deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/base_codec/PLACEHOLDER_base_codec_3.pth.tar.txt +++ /dev/null @@ -1,2 +0,0 @@ -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 deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/base_codec/PLACEHOLDER_base_codec_4.pth.tar.txt +++ /dev/null @@ -1,2 +0,0 @@ -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 deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage1/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1313.pth.tar b/checkpoints/ctaic/s1_det_instance/stage1/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage1/1/checkpoint_0.1313.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage1/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage1/2/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage1/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1334.pth.tar b/checkpoints/ctaic/s1_det_instance/stage1/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage1/2/checkpoint_0.1334.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage1/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage1/3/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage1/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1564.pth.tar b/checkpoints/ctaic/s1_det_instance/stage1/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage1/3/checkpoint_0.1564.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage1/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage1/4/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage1/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage1/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1245.pth.tar b/checkpoints/ctaic/s1_det_instance/stage1/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage1/4/checkpoint_0.1245.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage1/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage2/1/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage2/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1313.pth.tar b/checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage1/1/checkpoint_0.1313.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage2/2/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage2/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1334.pth.tar b/checkpoints/ctaic/s1_det_instance/stage2/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage1/2/checkpoint_0.1334.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage2/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage2/3/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage2/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1564.pth.tar b/checkpoints/ctaic/s1_det_instance/stage2/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage1/3/checkpoint_0.1564.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage2/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s1_det_instance/stage2/4/PLACEHOLDER b/checkpoints/ctaic/s1_det_instance/stage2/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s1_det_instance/stage2/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1245.pth.tar b/checkpoints/ctaic/s1_det_instance/stage2/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage1/4/checkpoint_0.1245.pth.tar rename to checkpoints/ctaic/s1_det_instance/stage2/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/1/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage1/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1313.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage1/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s3_det_pose/stage1/1/checkpoint_0.1313.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage1/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/2/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage1/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1334.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage1/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s3_det_pose/stage1/2/checkpoint_0.1334.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage1/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/3/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage1/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1564.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage1/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s3_det_pose/stage1/3/checkpoint_0.1564.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage1/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage1/4/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage1/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage1/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1245.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage1/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s3_det_pose/stage1/4/checkpoint_0.1245.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage1/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/1/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage2/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1313.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage2/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/detection/1/checkpoint_0.1313.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage2/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/2/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage2/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1334.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage2/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/detection/2/checkpoint_0.1334.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage2/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/3/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage2/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1564.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage2/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/detection/3/checkpoint_0.1564.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage2/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/4/PLACEHOLDER b/checkpoints/ctaic/s2_sem_panoptic/stage2/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s2_sem_panoptic/stage2/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1245.pth.tar b/checkpoints/ctaic/s2_sem_panoptic/stage2/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/detection/4/checkpoint_0.1245.pth.tar rename to checkpoints/ctaic/s2_sem_panoptic/stage2/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s3_det_pose/stage1/1/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage1/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1313.pth.tar b/checkpoints/ctaic/s3_det_pose/stage1/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/instance/1/checkpoint_0.1313.pth.tar rename to checkpoints/ctaic/s3_det_pose/stage1/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s3_det_pose/stage1/2/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage1/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1334.pth.tar b/checkpoints/ctaic/s3_det_pose/stage1/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/instance/2/checkpoint_0.1334.pth.tar rename to checkpoints/ctaic/s3_det_pose/stage1/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s3_det_pose/stage1/3/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage1/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1564.pth.tar b/checkpoints/ctaic/s3_det_pose/stage1/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/instance/3/checkpoint_0.1564.pth.tar rename to checkpoints/ctaic/s3_det_pose/stage1/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s3_det_pose/stage1/4/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage1/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage1/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.1245.pth.tar b/checkpoints/ctaic/s3_det_pose/stage1/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/taic/instance/4/checkpoint_0.1245.pth.tar rename to checkpoints/ctaic/s3_det_pose/stage1/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/ctaic/s3_det_pose/stage2/1/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0055.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_0.0055.pth.tar deleted file mode 100644 index 66f70b8678afcda9391313cd451d332d70d8e03b..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_0.0055.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a39ec8c02a28b7f41893e64e551e37b7d705b485e4394f04cc087580684f3b39 -size 72118222 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_best_loss.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/1/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/2/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0065.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/2/checkpoint_0.0065.pth.tar deleted file mode 100644 index 6f24417c77e3896f7d55af3df159797335dc14b2..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/2/checkpoint_0.0065.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d9270ac6300dd4cc6d9aee3f33d6f672cf0ef808a8d7840b6d1261309814546c -size 72118222 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/2/checkpoint_best_loss.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/2/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/2/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/3/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0204.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/3/checkpoint_0.0204.pth.tar deleted file mode 100644 index 4e21c0b1493e80e75b746462dd4c3eaf1cbbb34c..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/3/checkpoint_0.0204.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b07533a8a6e9dbc5fdcf0d9f23d84825c49469959d23e99409781f192b1a17e1 -size 72118158 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/3/checkpoint_best_loss.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/3/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/3/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/4/PLACEHOLDER b/checkpoints/ctaic/s3_det_pose/stage2/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0301.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/4/checkpoint_0.0301.pth.tar deleted file mode 100644 index e7d07f5d22e7c3ec0e582be4e45931dfa14b10a8..0000000000000000000000000000000000000000 --- a/checkpoints/ctaic/s3_det_pose/stage2/4/checkpoint_0.0301.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9e28094b13bbbe1ee167773106422a2c464aa503f7b29cf4946e85bbc2780b2a -size 72118158 diff --git a/checkpoints/ctaic/s3_det_pose/stage2/4/checkpoint_best_loss.pth.tar b/checkpoints/ctaic/s3_det_pose/stage2/4/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/ctaic/s3_det_pose/stage2/4/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/taic/detection/1/PLACEHOLDER b/checkpoints/taic/detection/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/detection/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..89ccb7bd842b4ded2b64de0cbaae3ed5767d641e --- /dev/null +++ b/checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d0a04a70771422dc827bde265ef4ef9b0f40cec25babfd640f2b5e803d2336c8 +size 72118222 diff --git a/checkpoints/taic/detection/2/PLACEHOLDER b/checkpoints/taic/detection/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/detection/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/detection/2/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..db609820b18ce90ecd6c8ee8868f4b30304525bd --- /dev/null +++ b/checkpoints/taic/detection/2/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:475653214e022e7aacdf9d3c711a6bde1185975eb1023dedb295faab87a790e5 +size 72118158 diff --git a/checkpoints/taic/detection/3/PLACEHOLDER b/checkpoints/taic/detection/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/detection/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/detection/3/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..d2f0312c8ed0f3e4970580c2b68eb32bcc34aa22 --- /dev/null +++ b/checkpoints/taic/detection/3/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2bf7ff870556e638938e081f8dce9dcadbfe3b72f108729810adf7d74390a7bf +size 72118158 diff --git a/checkpoints/taic/detection/4/PLACEHOLDER b/checkpoints/taic/detection/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/detection/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/detection/4/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..f9706f4ce967b7e0596e785d94641f631e08d865 --- /dev/null +++ b/checkpoints/taic/detection/4/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:991696cf49c62dbff8e244ab1ca5080aec880fa756177ba60baa11208c2e9d6c +size 72118158 diff --git a/checkpoints/taic/instance/1/PLACEHOLDER b/checkpoints/taic/instance/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/instance/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/instance/1/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..89ccb7bd842b4ded2b64de0cbaae3ed5767d641e --- /dev/null +++ b/checkpoints/taic/instance/1/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:d0a04a70771422dc827bde265ef4ef9b0f40cec25babfd640f2b5e803d2336c8 +size 72118222 diff --git a/checkpoints/taic/instance/2/PLACEHOLDER b/checkpoints/taic/instance/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/instance/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/instance/2/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..db609820b18ce90ecd6c8ee8868f4b30304525bd --- /dev/null +++ b/checkpoints/taic/instance/2/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:475653214e022e7aacdf9d3c711a6bde1185975eb1023dedb295faab87a790e5 +size 72118158 diff --git a/checkpoints/taic/instance/3/PLACEHOLDER b/checkpoints/taic/instance/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/instance/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/instance/3/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..d2f0312c8ed0f3e4970580c2b68eb32bcc34aa22 --- /dev/null +++ b/checkpoints/taic/instance/3/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:2bf7ff870556e638938e081f8dce9dcadbfe3b72f108729810adf7d74390a7bf +size 72118158 diff --git a/checkpoints/taic/instance/4/PLACEHOLDER b/checkpoints/taic/instance/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/instance/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_best_loss.pth.tar b/checkpoints/taic/instance/4/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..f9706f4ce967b7e0596e785d94641f631e08d865 --- /dev/null +++ b/checkpoints/taic/instance/4/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:991696cf49c62dbff8e244ab1ca5080aec880fa756177ba60baa11208c2e9d6c +size 72118158 diff --git a/checkpoints/taic/panoptic/1/PLACEHOLDER b/checkpoints/taic/panoptic/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0055.pth.tar b/checkpoints/taic/panoptic/1/checkpoint_0.0055.pth.tar deleted file mode 100644 index 66f70b8678afcda9391313cd451d332d70d8e03b..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/1/checkpoint_0.0055.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a39ec8c02a28b7f41893e64e551e37b7d705b485e4394f04cc087580684f3b39 -size 72118222 diff --git a/checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_0.0055.pth.tar b/checkpoints/taic/panoptic/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage2/1/checkpoint_0.0055.pth.tar rename to checkpoints/taic/panoptic/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/panoptic/2/PLACEHOLDER b/checkpoints/taic/panoptic/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0065.pth.tar b/checkpoints/taic/panoptic/2/checkpoint_0.0065.pth.tar deleted file mode 100644 index 6f24417c77e3896f7d55af3df159797335dc14b2..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/2/checkpoint_0.0065.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d9270ac6300dd4cc6d9aee3f33d6f672cf0ef808a8d7840b6d1261309814546c -size 72118222 diff --git a/checkpoints/ctaic/s1_det_instance/stage2/2/checkpoint_0.0065.pth.tar b/checkpoints/taic/panoptic/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage2/2/checkpoint_0.0065.pth.tar rename to checkpoints/taic/panoptic/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/panoptic/3/PLACEHOLDER b/checkpoints/taic/panoptic/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0204.pth.tar b/checkpoints/taic/panoptic/3/checkpoint_0.0204.pth.tar deleted file mode 100644 index 4e21c0b1493e80e75b746462dd4c3eaf1cbbb34c..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/3/checkpoint_0.0204.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b07533a8a6e9dbc5fdcf0d9f23d84825c49469959d23e99409781f192b1a17e1 -size 72118158 diff --git a/checkpoints/ctaic/s1_det_instance/stage2/3/checkpoint_0.0204.pth.tar b/checkpoints/taic/panoptic/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage2/3/checkpoint_0.0204.pth.tar rename to checkpoints/taic/panoptic/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/panoptic/4/PLACEHOLDER b/checkpoints/taic/panoptic/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0301.pth.tar b/checkpoints/taic/panoptic/4/checkpoint_0.0301.pth.tar deleted file mode 100644 index e7d07f5d22e7c3ec0e582be4e45931dfa14b10a8..0000000000000000000000000000000000000000 --- a/checkpoints/taic/panoptic/4/checkpoint_0.0301.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9e28094b13bbbe1ee167773106422a2c464aa503f7b29cf4946e85bbc2780b2a -size 72118158 diff --git a/checkpoints/ctaic/s1_det_instance/stage2/4/checkpoint_0.0301.pth.tar b/checkpoints/taic/panoptic/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s1_det_instance/stage2/4/checkpoint_0.0301.pth.tar rename to checkpoints/taic/panoptic/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/pose/1/PLACEHOLDER b/checkpoints/taic/pose/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0055.pth.tar b/checkpoints/taic/pose/1/checkpoint_0.0055.pth.tar deleted file mode 100644 index 66f70b8678afcda9391313cd451d332d70d8e03b..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/1/checkpoint_0.0055.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a39ec8c02a28b7f41893e64e551e37b7d705b485e4394f04cc087580684f3b39 -size 72118222 diff --git a/checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar b/checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/taic/pose/1/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/taic/pose/2/PLACEHOLDER b/checkpoints/taic/pose/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0065.pth.tar b/checkpoints/taic/pose/2/checkpoint_0.0065.pth.tar deleted file mode 100644 index 6f24417c77e3896f7d55af3df159797335dc14b2..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/2/checkpoint_0.0065.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d9270ac6300dd4cc6d9aee3f33d6f672cf0ef808a8d7840b6d1261309814546c -size 72118222 diff --git a/checkpoints/taic/pose/2/checkpoint_best_loss.pth.tar b/checkpoints/taic/pose/2/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/taic/pose/2/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/taic/pose/3/PLACEHOLDER b/checkpoints/taic/pose/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0204.pth.tar b/checkpoints/taic/pose/3/checkpoint_0.0204.pth.tar deleted file mode 100644 index 4e21c0b1493e80e75b746462dd4c3eaf1cbbb34c..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/3/checkpoint_0.0204.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b07533a8a6e9dbc5fdcf0d9f23d84825c49469959d23e99409781f192b1a17e1 -size 72118158 diff --git a/checkpoints/taic/pose/3/checkpoint_best_loss.pth.tar b/checkpoints/taic/pose/3/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/taic/pose/3/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/taic/pose/4/PLACEHOLDER b/checkpoints/taic/pose/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0301.pth.tar b/checkpoints/taic/pose/4/checkpoint_0.0301.pth.tar deleted file mode 100644 index e7d07f5d22e7c3ec0e582be4e45931dfa14b10a8..0000000000000000000000000000000000000000 --- a/checkpoints/taic/pose/4/checkpoint_0.0301.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9e28094b13bbbe1ee167773106422a2c464aa503f7b29cf4946e85bbc2780b2a -size 72118158 diff --git a/checkpoints/taic/pose/4/checkpoint_best_loss.pth.tar b/checkpoints/taic/pose/4/checkpoint_best_loss.pth.tar new file mode 100644 index 0000000000000000000000000000000000000000..41216ad122e03be2705b506b89f4598587f06525 --- /dev/null +++ b/checkpoints/taic/pose/4/checkpoint_best_loss.pth.tar @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:a8c869589438d291a34911e8dffe8a462e67d0587f94773cee66eb9ecd0a80b1 +size 70789902 diff --git a/checkpoints/taic/semantic/1/PLACEHOLDER b/checkpoints/taic/semantic/1/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/1/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0055.pth.tar b/checkpoints/taic/semantic/1/checkpoint_0.0055.pth.tar deleted file mode 100644 index 66f70b8678afcda9391313cd451d332d70d8e03b..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/1/checkpoint_0.0055.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:a39ec8c02a28b7f41893e64e551e37b7d705b485e4394f04cc087580684f3b39 -size 72118222 diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/1/checkpoint_0.0055.pth.tar b/checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage2/1/checkpoint_0.0055.pth.tar rename to checkpoints/taic/semantic/1/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/semantic/2/PLACEHOLDER b/checkpoints/taic/semantic/2/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/2/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0065.pth.tar b/checkpoints/taic/semantic/2/checkpoint_0.0065.pth.tar deleted file mode 100644 index 6f24417c77e3896f7d55af3df159797335dc14b2..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/2/checkpoint_0.0065.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:d9270ac6300dd4cc6d9aee3f33d6f672cf0ef808a8d7840b6d1261309814546c -size 72118222 diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/2/checkpoint_0.0065.pth.tar b/checkpoints/taic/semantic/2/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage2/2/checkpoint_0.0065.pth.tar rename to checkpoints/taic/semantic/2/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/semantic/3/PLACEHOLDER b/checkpoints/taic/semantic/3/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/3/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0204.pth.tar b/checkpoints/taic/semantic/3/checkpoint_0.0204.pth.tar deleted file mode 100644 index 4e21c0b1493e80e75b746462dd4c3eaf1cbbb34c..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/3/checkpoint_0.0204.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:b07533a8a6e9dbc5fdcf0d9f23d84825c49469959d23e99409781f192b1a17e1 -size 72118158 diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/3/checkpoint_0.0204.pth.tar b/checkpoints/taic/semantic/3/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage2/3/checkpoint_0.0204.pth.tar rename to checkpoints/taic/semantic/3/checkpoint_best_loss.pth.tar diff --git a/checkpoints/taic/semantic/4/PLACEHOLDER b/checkpoints/taic/semantic/4/PLACEHOLDER deleted file mode 100644 index c8768731085e17b2e27011b46e84e4423af2019f..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/4/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -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/checkpoint_0.0301.pth.tar b/checkpoints/taic/semantic/4/checkpoint_0.0301.pth.tar deleted file mode 100644 index e7d07f5d22e7c3ec0e582be4e45931dfa14b10a8..0000000000000000000000000000000000000000 --- a/checkpoints/taic/semantic/4/checkpoint_0.0301.pth.tar +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:9e28094b13bbbe1ee167773106422a2c464aa503f7b29cf4946e85bbc2780b2a -size 72118158 diff --git a/checkpoints/ctaic/s2_sem_panoptic/stage2/4/checkpoint_0.0301.pth.tar b/checkpoints/taic/semantic/4/checkpoint_best_loss.pth.tar similarity index 100% rename from checkpoints/ctaic/s2_sem_panoptic/stage2/4/checkpoint_0.0301.pth.tar rename to checkpoints/taic/semantic/4/checkpoint_best_loss.pth.tar diff --git a/checkpoints/task_networks/detection/PLACEHOLDER b/checkpoints/task_networks/detection/PLACEHOLDER deleted file mode 100644 index e76239419c8d8f9ae5fb43310ded28df1ac43155..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/detection/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -PLACEHOLDER: put the official pretrained task-network weight here as model.pth -Also set task_config in configs/eval/*.yaml to the matching OpenMMLab config. diff --git a/checkpoints/task_networks/instance/PLACEHOLDER b/checkpoints/task_networks/instance/PLACEHOLDER deleted file mode 100644 index e76239419c8d8f9ae5fb43310ded28df1ac43155..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/instance/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -PLACEHOLDER: put the official pretrained task-network weight here as model.pth -Also set task_config in configs/eval/*.yaml to the matching OpenMMLab config. diff --git a/checkpoints/task_networks/panoptic/PLACEHOLDER b/checkpoints/task_networks/panoptic/PLACEHOLDER deleted file mode 100644 index e76239419c8d8f9ae5fb43310ded28df1ac43155..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/panoptic/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -PLACEHOLDER: put the official pretrained task-network weight here as model.pth -Also set task_config in configs/eval/*.yaml to the matching OpenMMLab config. diff --git a/checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py b/checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py deleted file mode 100644 index 78c680147bb0fc48be52084a09f93424a5340430..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py +++ /dev/null @@ -1,778 +0,0 @@ -auto_scale_lr = dict(base_batch_size=16, enable=False) -backbone_embed_multi = dict(decay_mult=0.0, lr_mult=0.1) -backbone_norm_multi = dict(decay_mult=0.0, lr_mult=0.1) -backend_args = None -batch_augments = [ - dict( - img_pad_value=0, - mask_pad_value=0, - pad_mask=True, - pad_seg=True, - seg_pad_value=255, - size=( - 1024, - 1024, - ), - type='BatchFixedSizePad'), -] -custom_keys = dict({ - 'absolute_pos_embed': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone': - dict(decay_mult=1.0, lr_mult=0.1), - 'backbone.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.patch_embed.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.0.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.0.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.0.downsample.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.1.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.1.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.1.downsample.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.10.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.11.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.12.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.13.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.14.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.15.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.16.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.17.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.2.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.3.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.4.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.5.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.6.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.7.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.8.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.9.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.downsample.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.3.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.3.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'level_embed': - dict(decay_mult=0.0, lr_mult=1.0), - 'query_embed': - dict(decay_mult=0.0, lr_mult=1.0), - 'query_feat': - dict(decay_mult=0.0, lr_mult=1.0), - 'relative_position_bias_table': - dict(decay_mult=0.0, lr_mult=0.1) -}) -data_preprocessor = dict( - batch_augments=[ - dict( - img_pad_value=0, - mask_pad_value=0, - pad_mask=True, - pad_seg=True, - seg_pad_value=255, - size=( - 1024, - 1024, - ), - type='BatchFixedSizePad'), - ], - bgr_to_rgb=True, - mask_pad_value=0, - mean=[ - 123.675, - 116.28, - 103.53, - ], - pad_mask=True, - pad_seg=True, - pad_size_divisor=32, - seg_pad_value=255, - std=[ - 58.395, - 57.12, - 57.375, - ], - type='DetDataPreprocessor') -data_root = 'data/coco/' -dataset_type = 'CocoPanopticDataset' -default_hooks = dict( - checkpoint=dict( - by_epoch=False, - interval=5000, - max_keep_ckpts=3, - save_last=True, - type='CheckpointHook'), - logger=dict(interval=50, type='LoggerHook'), - param_scheduler=dict(type='ParamSchedulerHook'), - sampler_seed=dict(type='DistSamplerSeedHook'), - timer=dict(type='IterTimerHook'), - visualization=dict(type='DetVisualizationHook')) -default_scope = 'mmdet' -depths = [ - 2, - 2, - 18, - 2, -] -dynamic_intervals = [ - ( - 365001, - 368750, - ), -] -embed_multi = dict(decay_mult=0.0, lr_mult=1.0) -env_cfg = dict( - cudnn_benchmark=False, - dist_cfg=dict(backend='nccl'), - mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0)) -image_size = ( - 1024, - 1024, -) -interval = 5000 -load_from = None -log_level = 'INFO' -log_processor = dict(by_epoch=False, type='LogProcessor', window_size=50) -max_iters = 368750 -model = dict( - backbone=dict( - attn_drop_rate=0.0, - convert_weights=True, - depths=[ - 2, - 2, - 18, - 2, - ], - drop_path_rate=0.3, - drop_rate=0.0, - embed_dims=128, - frozen_stages=-1, - init_cfg=dict( - checkpoint= - 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth', - type='Pretrained'), - mlp_ratio=4, - num_heads=[ - 4, - 8, - 16, - 32, - ], - out_indices=( - 0, - 1, - 2, - 3, - ), - patch_norm=True, - pretrain_img_size=384, - qk_scale=None, - qkv_bias=True, - type='SwinTransformer', - window_size=12, - with_cp=False), - data_preprocessor=dict( - batch_augments=[ - dict( - img_pad_value=0, - mask_pad_value=0, - pad_mask=True, - pad_seg=True, - seg_pad_value=255, - size=( - 1024, - 1024, - ), - type='BatchFixedSizePad'), - ], - bgr_to_rgb=True, - mask_pad_value=0, - mean=[ - 123.675, - 116.28, - 103.53, - ], - pad_mask=True, - pad_seg=True, - pad_size_divisor=32, - seg_pad_value=255, - std=[ - 58.395, - 57.12, - 57.375, - ], - type='DetDataPreprocessor'), - init_cfg=None, - panoptic_fusion_head=dict( - init_cfg=None, - loss_panoptic=None, - num_stuff_classes=53, - num_things_classes=80, - type='MaskFormerFusionHead'), - panoptic_head=dict( - enforce_decoder_input_project=False, - feat_channels=256, - in_channels=[ - 128, - 256, - 512, - 1024, - ], - loss_cls=dict( - class_weight=[ - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 1.0, - 0.1, - ], - loss_weight=2.0, - reduction='mean', - type='CrossEntropyLoss', - use_sigmoid=False), - loss_dice=dict( - activate=True, - eps=1.0, - loss_weight=5.0, - naive_dice=True, - reduction='mean', - type='DiceLoss', - use_sigmoid=True), - loss_mask=dict( - loss_weight=5.0, - reduction='mean', - type='CrossEntropyLoss', - use_sigmoid=True), - num_queries=100, - num_stuff_classes=53, - num_things_classes=80, - num_transformer_feat_level=3, - out_channels=256, - pixel_decoder=dict( - act_cfg=dict(type='ReLU'), - encoder=dict( - layer_cfg=dict( - ffn_cfg=dict( - act_cfg=dict(inplace=True, type='ReLU'), - embed_dims=256, - feedforward_channels=1024, - ffn_drop=0.0, - num_fcs=2), - self_attn_cfg=dict( - batch_first=True, - dropout=0.0, - embed_dims=256, - num_heads=8, - num_levels=3, - num_points=4)), - num_layers=6), - norm_cfg=dict(num_groups=32, type='GN'), - num_outs=3, - positional_encoding=dict(normalize=True, num_feats=128), - type='MSDeformAttnPixelDecoder'), - positional_encoding=dict(normalize=True, num_feats=128), - strides=[ - 4, - 8, - 16, - 32, - ], - transformer_decoder=dict( - init_cfg=None, - layer_cfg=dict( - cross_attn_cfg=dict( - batch_first=True, dropout=0.0, embed_dims=256, - num_heads=8), - ffn_cfg=dict( - act_cfg=dict(inplace=True, type='ReLU'), - embed_dims=256, - feedforward_channels=2048, - ffn_drop=0.0, - num_fcs=2), - self_attn_cfg=dict( - batch_first=True, dropout=0.0, embed_dims=256, - num_heads=8)), - num_layers=9, - return_intermediate=True), - type='Mask2FormerHead'), - test_cfg=dict( - filter_low_score=True, - instance_on=True, - iou_thr=0.8, - max_per_image=100, - panoptic_on=True, - semantic_on=False), - train_cfg=dict( - assigner=dict( - match_costs=[ - dict(type='ClassificationCost', weight=2.0), - dict( - type='CrossEntropyLossCost', use_sigmoid=True, weight=5.0), - dict(eps=1.0, pred_act=True, type='DiceCost', weight=5.0), - ], - type='HungarianAssigner'), - importance_sample_ratio=0.75, - num_points=12544, - oversample_ratio=3.0, - sampler=dict(type='MaskPseudoSampler')), - type='Mask2Former') -num_classes = 133 -num_stuff_classes = 53 -num_things_classes = 80 -optim_wrapper = dict( - clip_grad=dict(max_norm=0.01, norm_type=2), - optimizer=dict( - betas=( - 0.9, - 0.999, - ), - eps=1e-08, - lr=0.0001, - type='AdamW', - weight_decay=0.05), - paramwise_cfg=dict( - custom_keys=dict({ - 'absolute_pos_embed': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone': - dict(decay_mult=1.0, lr_mult=0.1), - 'backbone.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.patch_embed.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.0.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.0.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.0.downsample.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.1.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.1.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.1.downsample.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.10.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.11.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.12.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.13.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.14.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.15.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.16.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.17.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.2.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.3.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.4.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.5.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.6.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.7.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.8.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.blocks.9.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.2.downsample.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.3.blocks.0.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'backbone.stages.3.blocks.1.norm': - dict(decay_mult=0.0, lr_mult=0.1), - 'level_embed': - dict(decay_mult=0.0, lr_mult=1.0), - 'query_embed': - dict(decay_mult=0.0, lr_mult=1.0), - 'query_feat': - dict(decay_mult=0.0, lr_mult=1.0), - 'relative_position_bias_table': - dict(decay_mult=0.0, lr_mult=0.1) - }), - norm_decay_mult=0.0), - type='OptimWrapper') -param_scheduler = dict( - begin=0, - by_epoch=False, - end=368750, - gamma=0.1, - milestones=[ - 327778, - 355092, - ], - type='MultiStepLR') -pretrained = 'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth' -resume = False -test_cfg = dict(type='TestLoop') -test_dataloader = dict( - batch_size=1, - dataset=dict( - ann_file='annotations/panoptic_val2017.json', - backend_args=None, - data_prefix=dict(img='val2017/', seg='annotations/panoptic_val2017/'), - data_root='data/coco/', - pipeline=[ - dict(backend_args=None, type='LoadImageFromFile'), - dict(keep_ratio=True, scale=( - 1333, - 800, - ), type='Resize'), - dict(backend_args=None, type='LoadPanopticAnnotations'), - dict( - meta_keys=( - 'img_id', - 'img_path', - 'ori_shape', - 'img_shape', - 'scale_factor', - ), - type='PackDetInputs'), - ], - test_mode=True, - type='CocoPanopticDataset'), - drop_last=False, - num_workers=2, - persistent_workers=True, - sampler=dict(shuffle=False, type='DefaultSampler')) -test_evaluator = [ - dict( - ann_file='data/coco/annotations/panoptic_val2017.json', - backend_args=None, - seg_prefix='data/coco/annotations/panoptic_val2017/', - type='CocoPanopticMetric'), - dict( - ann_file='data/coco/annotations/instances_val2017.json', - backend_args=None, - metric=[ - 'bbox', - 'segm', - ], - type='CocoMetric'), -] -test_pipeline = [ - dict(backend_args=None, type='LoadImageFromFile'), - dict(keep_ratio=True, scale=( - 1333, - 800, - ), type='Resize'), - dict(backend_args=None, type='LoadPanopticAnnotations'), - dict( - meta_keys=( - 'img_id', - 'img_path', - 'ori_shape', - 'img_shape', - 'scale_factor', - ), - type='PackDetInputs'), -] -train_cfg = dict( - dynamic_intervals=[ - ( - 365001, - 368750, - ), - ], - max_iters=368750, - type='IterBasedTrainLoop', - val_interval=5000) -train_dataloader = dict( - batch_sampler=dict(type='AspectRatioBatchSampler'), - batch_size=2, - dataset=dict( - ann_file='annotations/panoptic_train2017.json', - backend_args=None, - data_prefix=dict( - img='train2017/', seg='annotations/panoptic_train2017/'), - data_root='data/coco/', - filter_cfg=dict(filter_empty_gt=True, min_size=32), - pipeline=[ - dict(backend_args=None, to_float32=True, type='LoadImageFromFile'), - dict( - backend_args=None, - type='LoadPanopticAnnotations', - with_bbox=True, - with_mask=True, - with_seg=True), - dict(prob=0.5, type='RandomFlip'), - dict( - keep_ratio=True, - ratio_range=( - 0.1, - 2.0, - ), - scale=( - 1024, - 1024, - ), - type='RandomResize'), - dict( - allow_negative_crop=True, - crop_size=( - 1024, - 1024, - ), - crop_type='absolute', - recompute_bbox=True, - type='RandomCrop'), - dict(type='PackDetInputs'), - ], - type='CocoPanopticDataset'), - num_workers=2, - persistent_workers=True, - sampler=dict(shuffle=True, type='DefaultSampler')) -train_pipeline = [ - dict(backend_args=None, to_float32=True, type='LoadImageFromFile'), - dict( - backend_args=None, - type='LoadPanopticAnnotations', - with_bbox=True, - with_mask=True, - with_seg=True), - dict(prob=0.5, type='RandomFlip'), - dict( - keep_ratio=True, - ratio_range=( - 0.1, - 2.0, - ), - scale=( - 1024, - 1024, - ), - type='RandomResize'), - dict( - allow_negative_crop=True, - crop_size=( - 1024, - 1024, - ), - crop_type='absolute', - recompute_bbox=True, - type='RandomCrop'), - dict(type='PackDetInputs'), -] -val_cfg = dict(type='ValLoop') -val_dataloader = dict( - batch_size=1, - dataset=dict( - ann_file='annotations/panoptic_val2017.json', - backend_args=None, - data_prefix=dict(img='val2017/', seg='annotations/panoptic_val2017/'), - data_root='data/coco/', - pipeline=[ - dict(backend_args=None, type='LoadImageFromFile'), - dict(keep_ratio=True, scale=( - 1333, - 800, - ), type='Resize'), - dict(backend_args=None, type='LoadPanopticAnnotations'), - dict( - meta_keys=( - 'img_id', - 'img_path', - 'ori_shape', - 'img_shape', - 'scale_factor', - ), - type='PackDetInputs'), - ], - test_mode=True, - type='CocoPanopticDataset'), - drop_last=False, - num_workers=2, - persistent_workers=True, - sampler=dict(shuffle=False, type='DefaultSampler')) -val_evaluator = [ - dict( - ann_file='data/coco/annotations/panoptic_val2017.json', - backend_args=None, - seg_prefix='data/coco/annotations/panoptic_val2017/', - type='CocoPanopticMetric'), - dict( - ann_file='data/coco/annotations/instances_val2017.json', - backend_args=None, - metric=[ - 'bbox', - 'segm', - ], - type='CocoMetric'), -] -vis_backends = [ - dict(type='LocalVisBackend'), -] -visualizer = dict( - name='visualizer', - type='DetLocalVisualizer', - vis_backends=[ - dict(type='LocalVisBackend'), - ]) diff --git a/checkpoints/task_networks/pose/PLACEHOLDER b/checkpoints/task_networks/pose/PLACEHOLDER deleted file mode 100644 index e76239419c8d8f9ae5fb43310ded28df1ac43155..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/pose/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -PLACEHOLDER: put the official pretrained task-network weight here as model.pth -Also set task_config in configs/eval/*.yaml to the matching OpenMMLab config. diff --git a/checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py b/checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py deleted file mode 100644 index d0e7768d6dd3a3ce7a78f5ea59fab3b9ef2fe1c3..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py +++ /dev/null @@ -1,350 +0,0 @@ -auto_scale_lr = dict(base_batch_size=192) -backend_args = dict(backend='local') -codec = dict( - decode_center_shift=0.5, - decode_keypoint_order=[ - 0, - 1, - 2, - 3, - 4, - 5, - 6, - 11, - 12, - 7, - 8, - 9, - 10, - 13, - 14, - 15, - 16, - ], - decode_max_instances=30, - decode_topk=30, - heatmap_size=( - 128, - 128, - ), - input_size=( - 512, - 512, - ), - sigma=2, - type='AssociativeEmbedding') -custom_hooks = [ - dict(type='SyncBuffersHook'), -] -data_mode = 'bottomup' -data_root = 'data/coco/' -dataset_type = 'CocoDataset' -default_hooks = dict( - badcase=dict( - badcase_thr=5, - enable=False, - metric_type='loss', - out_dir='badcase', - type='BadCaseAnalysisHook'), - checkpoint=dict( - interval=50, - rule='greater', - save_best='coco/AP', - type='CheckpointHook'), - logger=dict(interval=50, type='LoggerHook'), - param_scheduler=dict(type='ParamSchedulerHook'), - sampler_seed=dict(type='DistSamplerSeedHook'), - timer=dict(type='IterTimerHook'), - visualization=dict(enable=False, type='PoseVisualizationHook')) -default_scope = 'mmpose' -env_cfg = dict( - cudnn_benchmark=False, - dist_cfg=dict(backend='nccl'), - mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0)) -load_from = None -log_level = 'INFO' -log_processor = dict( - by_epoch=True, num_digits=6, type='LogProcessor', window_size=50) -model = dict( - backbone=dict( - extra=dict( - stage1=dict( - block='BOTTLENECK', - num_blocks=(4, ), - num_branches=1, - num_channels=(64, ), - num_modules=1), - stage2=dict( - block='BASIC', - num_blocks=( - 4, - 4, - ), - num_branches=2, - num_channels=( - 32, - 64, - ), - num_modules=1), - stage3=dict( - block='BASIC', - num_blocks=( - 4, - 4, - 4, - ), - num_branches=3, - num_channels=( - 32, - 64, - 128, - ), - num_modules=4), - stage4=dict( - block='BASIC', - num_blocks=( - 4, - 4, - 4, - 4, - ), - num_branches=4, - num_channels=( - 32, - 64, - 128, - 256, - ), - num_modules=3)), - in_channels=3, - init_cfg=dict( - checkpoint= - 'https://download.openmmlab.com/mmpose/pretrain_models/hrnet_w32-36af842e.pth', - type='Pretrained'), - type='HRNet'), - data_preprocessor=dict( - bgr_to_rgb=True, - mean=[ - 123.675, - 116.28, - 103.53, - ], - std=[ - 58.395, - 57.12, - 57.375, - ], - type='PoseDataPreprocessor'), - head=dict( - decoder=dict( - decode_center_shift=0.5, - decode_keypoint_order=[ - 0, - 1, - 2, - 3, - 4, - 5, - 6, - 11, - 12, - 7, - 8, - 9, - 10, - 13, - 14, - 15, - 16, - ], - decode_max_instances=30, - decode_topk=30, - heatmap_size=( - 512, - 512, - ), - input_size=( - 512, - 512, - ), - sigma=2, - type='AssociativeEmbedding'), - deconv_out_channels=None, - in_channels=32, - keypoint_loss=dict(type='KeypointMSELoss', use_target_weight=True), - num_keypoints=17, - tag_dim=1, - tag_loss=dict(loss_weight=0.001, type='AssociativeEmbeddingLoss'), - tag_per_keypoint=True, - type='AssociativeEmbeddingHead'), - test_cfg=dict( - align_corners=False, - flip_test=True, - multiscale_test=False, - restore_heatmap_size=True, - shift_heatmap=False), - type='BottomupPoseEstimator') -optim_wrapper = dict(optimizer=dict(lr=0.0015, type='Adam')) -param_scheduler = [ - dict( - begin=0, by_epoch=False, end=500, start_factor=0.001, type='LinearLR'), - dict( - begin=0, - by_epoch=True, - end=300, - gamma=0.1, - milestones=[ - 200, - 260, - ], - type='MultiStepLR'), -] -resume = False -test_cfg = dict() -test_dataloader = dict( - batch_size=1, - dataset=dict( - ann_file='annotations/person_keypoints_val2017.json', - data_mode='bottomup', - data_prefix=dict(img='val2017/'), - data_root='data/coco/', - pipeline=[ - dict(type='LoadImage'), - dict( - input_size=( - 512, - 512, - ), - resize_mode='expand', - size_factor=64, - type='BottomupResize'), - dict( - meta_keys=( - 'id', - 'img_id', - 'img_path', - 'crowd_index', - 'ori_shape', - 'img_shape', - 'input_size', - 'input_center', - 'input_scale', - 'flip', - 'flip_direction', - 'flip_indices', - 'raw_ann_info', - 'skeleton_links', - ), - type='PackPoseInputs'), - ], - test_mode=True, - type='CocoDataset'), - drop_last=False, - num_workers=2, - persistent_workers=True, - sampler=dict(round_up=False, shuffle=False, type='DefaultSampler')) -test_evaluator = dict( - ann_file='data/coco/annotations/person_keypoints_val2017.json', - nms_mode='none', - score_mode='bbox', - type='CocoMetric') -train_cfg = dict(by_epoch=True, max_epochs=300, val_interval=10) -train_dataloader = dict( - batch_size=24, - dataset=dict( - ann_file='annotations/person_keypoints_train2017.json', - data_mode='bottomup', - data_prefix=dict(img='train2017/'), - data_root='data/coco/', - pipeline=[], - type='CocoDataset'), - num_workers=2, - persistent_workers=True, - sampler=dict(shuffle=True, type='DefaultSampler')) -train_pipeline = [] -val_cfg = dict() -val_dataloader = dict( - batch_size=1, - dataset=dict( - ann_file='annotations/person_keypoints_val2017.json', - data_mode='bottomup', - data_prefix=dict(img='val2017/'), - data_root='data/coco/', - pipeline=[ - dict(type='LoadImage'), - dict( - input_size=( - 512, - 512, - ), - resize_mode='expand', - size_factor=64, - type='BottomupResize'), - dict( - meta_keys=( - 'id', - 'img_id', - 'img_path', - 'crowd_index', - 'ori_shape', - 'img_shape', - 'input_size', - 'input_center', - 'input_scale', - 'flip', - 'flip_direction', - 'flip_indices', - 'raw_ann_info', - 'skeleton_links', - ), - type='PackPoseInputs'), - ], - test_mode=True, - type='CocoDataset'), - drop_last=False, - num_workers=2, - persistent_workers=True, - sampler=dict(round_up=False, shuffle=False, type='DefaultSampler')) -val_evaluator = dict( - ann_file='data/coco/annotations/person_keypoints_val2017.json', - nms_mode='none', - score_mode='bbox', - type='CocoMetric') -val_pipeline = [ - dict(type='LoadImage'), - dict( - input_size=( - 512, - 512, - ), - resize_mode='expand', - size_factor=64, - type='BottomupResize'), - dict( - meta_keys=( - 'id', - 'img_id', - 'img_path', - 'crowd_index', - 'ori_shape', - 'img_shape', - 'input_size', - 'input_center', - 'input_scale', - 'flip', - 'flip_direction', - 'flip_indices', - 'raw_ann_info', - 'skeleton_links', - ), - type='PackPoseInputs'), -] -vis_backends = [ - dict(type='LocalVisBackend'), -] -visualizer = dict( - name='visualizer', - type='PoseLocalVisualizer', - vis_backends=[ - dict(type='LocalVisBackend'), - ]) diff --git a/checkpoints/task_networks/semantic/PLACEHOLDER b/checkpoints/task_networks/semantic/PLACEHOLDER deleted file mode 100644 index e76239419c8d8f9ae5fb43310ded28df1ac43155..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/semantic/PLACEHOLDER +++ /dev/null @@ -1,2 +0,0 @@ -PLACEHOLDER: put the official pretrained task-network weight here as model.pth -Also set task_config in configs/eval/*.yaml to the matching OpenMMLab config. diff --git a/checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py b/checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py deleted file mode 100644 index 3a98429858814afce982a416861cccf147920c7c..0000000000000000000000000000000000000000 --- a/checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py +++ /dev/null @@ -1,336 +0,0 @@ -backbone_norm_cfg = dict(requires_grad=True, type='LN') -checkpoint_file = 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/swin/swin_base_patch4_window7_224_22k_20220317-4f79f7c0.pth' -crop_size = ( - 512, - 512, -) -data_preprocessor = dict( - bgr_to_rgb=True, - mean=[ - 123.675, - 116.28, - 103.53, - ], - pad_val=0, - seg_pad_val=255, - size=( - 512, - 512, - ), - std=[ - 58.395, - 57.12, - 57.375, - ], - type='SegDataPreProcessor') -data_root = 'data/ade/ADEChallengeData2016' -dataset_type = 'ADE20KDataset' -default_hooks = dict( - checkpoint=dict(by_epoch=False, interval=16000, type='CheckpointHook'), - logger=dict(interval=50, log_metric_by_epoch=False, type='LoggerHook'), - param_scheduler=dict(type='ParamSchedulerHook'), - sampler_seed=dict(type='DistSamplerSeedHook'), - timer=dict(type='IterTimerHook'), - visualization=dict(type='SegVisualizationHook')) -default_scope = 'mmseg' -env_cfg = dict( - cudnn_benchmark=True, - dist_cfg=dict(backend='nccl'), - mp_cfg=dict(mp_start_method='fork', opencv_num_threads=0)) -img_ratios = [ - 0.5, - 0.75, - 1.0, - 1.25, - 1.5, - 1.75, -] -load_from = None -log_level = 'INFO' -log_processor = dict(by_epoch=False) -model = dict( - auxiliary_head=dict( - align_corners=False, - channels=256, - concat_input=False, - dropout_ratio=0.1, - in_channels=512, - in_index=2, - loss_decode=dict( - loss_weight=0.4, type='CrossEntropyLoss', use_sigmoid=False), - norm_cfg=dict(requires_grad=True, type='SyncBN'), - num_classes=150, - num_convs=1, - type='FCNHead'), - backbone=dict( - act_cfg=dict(type='GELU'), - attn_drop_rate=0.0, - depths=[ - 2, - 2, - 18, - 2, - ], - drop_path_rate=0.3, - drop_rate=0.0, - embed_dims=128, - init_cfg=dict( - checkpoint= - 'https://download.openmmlab.com/mmsegmentation/v0.5/pretrain/swin/swin_base_patch4_window7_224_22k_20220317-4f79f7c0.pth', - type='Pretrained'), - mlp_ratio=4, - norm_cfg=dict(requires_grad=True, type='LN'), - num_heads=[ - 4, - 8, - 16, - 32, - ], - out_indices=( - 0, - 1, - 2, - 3, - ), - patch_norm=True, - patch_size=4, - pretrain_img_size=224, - qk_scale=None, - qkv_bias=True, - strides=( - 4, - 2, - 2, - 2, - ), - type='SwinTransformer', - use_abs_pos_embed=False, - window_size=7), - data_preprocessor=dict( - bgr_to_rgb=True, - mean=[ - 123.675, - 116.28, - 103.53, - ], - pad_val=0, - seg_pad_val=255, - size=( - 512, - 512, - ), - std=[ - 58.395, - 57.12, - 57.375, - ], - type='SegDataPreProcessor'), - decode_head=dict( - align_corners=False, - channels=512, - dropout_ratio=0.1, - in_channels=[ - 128, - 256, - 512, - 1024, - ], - in_index=[ - 0, - 1, - 2, - 3, - ], - loss_decode=dict( - loss_weight=1.0, type='CrossEntropyLoss', use_sigmoid=False), - norm_cfg=dict(requires_grad=True, type='SyncBN'), - num_classes=150, - pool_scales=( - 1, - 2, - 3, - 6, - ), - type='UPerHead'), - pretrained=None, - test_cfg=dict(mode='whole'), - train_cfg=dict(), - type='EncoderDecoder') -norm_cfg = dict(requires_grad=True, type='SyncBN') -optim_wrapper = dict( - optimizer=dict( - betas=( - 0.9, - 0.999, - ), lr=6e-05, type='AdamW', weight_decay=0.01), - paramwise_cfg=dict( - custom_keys=dict( - absolute_pos_embed=dict(decay_mult=0.0), - norm=dict(decay_mult=0.0), - relative_position_bias_table=dict(decay_mult=0.0))), - type='OptimWrapper') -optimizer = dict(lr=0.01, momentum=0.9, type='SGD', weight_decay=0.0005) -param_scheduler = [ - dict( - begin=0, by_epoch=False, end=1500, start_factor=1e-06, - type='LinearLR'), - dict( - begin=1500, - by_epoch=False, - end=160000, - eta_min=0.0, - power=1.0, - type='PolyLR'), -] -resume = False -test_cfg = dict(type='TestLoop') -test_dataloader = dict( - batch_size=1, - dataset=dict( - data_prefix=dict( - img_path='images/validation', - seg_map_path='annotations/validation'), - data_root='data/ade/ADEChallengeData2016', - pipeline=[ - dict(type='LoadImageFromFile'), - dict(keep_ratio=True, scale=( - 2048, - 512, - ), type='Resize'), - dict(reduce_zero_label=True, type='LoadAnnotations'), - dict(type='PackSegInputs'), - ], - type='ADE20KDataset'), - num_workers=4, - persistent_workers=True, - sampler=dict(shuffle=False, type='DefaultSampler')) -test_evaluator = dict( - iou_metrics=[ - 'mIoU', - ], type='IoUMetric') -test_pipeline = [ - dict(type='LoadImageFromFile'), - dict(keep_ratio=True, scale=( - 2048, - 512, - ), type='Resize'), - dict(reduce_zero_label=True, type='LoadAnnotations'), - dict(type='PackSegInputs'), -] -train_cfg = dict( - max_iters=160000, type='IterBasedTrainLoop', val_interval=16000) -train_dataloader = dict( - batch_size=2, - dataset=dict( - data_prefix=dict( - img_path='images/training', seg_map_path='annotations/training'), - data_root='data/ade/ADEChallengeData2016', - pipeline=[ - dict(type='LoadImageFromFile'), - dict(reduce_zero_label=True, type='LoadAnnotations'), - dict( - keep_ratio=True, - ratio_range=( - 0.5, - 2.0, - ), - scale=( - 2048, - 512, - ), - type='RandomResize'), - dict( - cat_max_ratio=0.75, crop_size=( - 512, - 512, - ), type='RandomCrop'), - dict(prob=0.5, type='RandomFlip'), - dict(type='PhotoMetricDistortion'), - dict(type='PackSegInputs'), - ], - type='ADE20KDataset'), - num_workers=4, - persistent_workers=True, - sampler=dict(shuffle=True, type='InfiniteSampler')) -train_pipeline = [ - dict(type='LoadImageFromFile'), - dict(reduce_zero_label=True, type='LoadAnnotations'), - dict( - keep_ratio=True, - ratio_range=( - 0.5, - 2.0, - ), - scale=( - 2048, - 512, - ), - type='RandomResize'), - dict(cat_max_ratio=0.75, crop_size=( - 512, - 512, - ), type='RandomCrop'), - dict(prob=0.5, type='RandomFlip'), - dict(type='PhotoMetricDistortion'), - dict(type='PackSegInputs'), -] -tta_model = dict(type='SegTTAModel') -tta_pipeline = [ - dict(backend_args=None, type='LoadImageFromFile'), - dict( - transforms=[ - [ - dict(keep_ratio=True, scale_factor=0.5, type='Resize'), - dict(keep_ratio=True, scale_factor=0.75, type='Resize'), - dict(keep_ratio=True, scale_factor=1.0, type='Resize'), - dict(keep_ratio=True, scale_factor=1.25, type='Resize'), - dict(keep_ratio=True, scale_factor=1.5, type='Resize'), - dict(keep_ratio=True, scale_factor=1.75, type='Resize'), - ], - [ - dict(direction='horizontal', prob=0.0, type='RandomFlip'), - dict(direction='horizontal', prob=1.0, type='RandomFlip'), - ], - [ - dict(type='LoadAnnotations'), - ], - [ - dict(type='PackSegInputs'), - ], - ], - type='TestTimeAug'), -] -val_cfg = dict(type='ValLoop') -val_dataloader = dict( - batch_size=1, - dataset=dict( - data_prefix=dict( - img_path='images/validation', - seg_map_path='annotations/validation'), - data_root='data/ade/ADEChallengeData2016', - pipeline=[ - dict(type='LoadImageFromFile'), - dict(keep_ratio=True, scale=( - 2048, - 512, - ), type='Resize'), - dict(reduce_zero_label=True, type='LoadAnnotations'), - dict(type='PackSegInputs'), - ], - type='ADE20KDataset'), - num_workers=4, - persistent_workers=True, - sampler=dict(shuffle=False, type='DefaultSampler')) -val_evaluator = dict( - iou_metrics=[ - 'mIoU', - ], type='IoUMetric') -vis_backends = [ - dict(type='LocalVisBackend'), -] -visualizer = dict( - name='visualizer', - type='SegLocalVisualizer', - vis_backends=[ - dict(type='LocalVisBackend'), - ]) diff --git a/configs/ctaic/s1_det_instance.yaml b/configs/ctaic/s1_det_instance.yaml deleted file mode 100644 index f57be95402bd9db2ab616d536c2f6d57da8f1674..0000000000000000000000000000000000000000 --- a/configs/ctaic/s1_det_instance.yaml +++ /dev/null @@ -1,31 +0,0 @@ -# Scenario s1: Object Detection (base) + Instance Segmentation (extension) -root: "logs" -exp_name: "ctaic_s1" -scenario: "s1" -dataset_path: "/data/Dataset/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 -out_channels: 128 # extension = Cascade Mask R-CNN + Swin-B -align_mode: "fpn" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth" diff --git a/configs/ctaic/s2_sem_panoptic.yaml b/configs/ctaic/s2_sem_panoptic.yaml deleted file mode 100644 index d044b33d37fba67b188fe983f0accdf3ca9e0c8d..0000000000000000000000000000000000000000 --- a/configs/ctaic/s2_sem_panoptic.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# Scenario s2: Semantic Segmentation (base) + Panoptic Segmentation (extension) -root: "logs" -exp_name: "ctaic_s2" -scenario: "s2" -dataset_path: "/data/Dataset/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 - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py" -task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth" diff --git a/configs/ctaic/s3_det_pose.yaml b/configs/ctaic/s3_det_pose.yaml deleted file mode 100644 index be828dcf5013b5300e1d65f73a9ffc859c0e5c7c..0000000000000000000000000000000000000000 --- a/configs/ctaic/s3_det_pose.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# Scenario s3: Object Detection (base) + Pose Estimation (extension) -root: "logs" -exp_name: "ctaic_s3" -scenario: "s3" -dataset_path: "/data/Dataset/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" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py" -task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth" diff --git a/configs/eval/ctaic_s1.yaml b/configs/eval/ctaic_s1.yaml deleted file mode 100644 index 8e0dfed0840c30278e40785f4075ae5a06db9fe5..0000000000000000000000000000000000000000 --- a/configs/eval/ctaic_s1.yaml +++ /dev/null @@ -1,19 +0,0 @@ -scenario: "s1" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/instances_val2017.json" -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 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth" diff --git a/configs/eval/ctaic_s2.yaml b/configs/eval/ctaic_s2.yaml deleted file mode 100644 index 63e51d2e36b3477d1000b1c667b72ad84edcbf24..0000000000000000000000000000000000000000 --- a/configs/eval/ctaic_s2.yaml +++ /dev/null @@ -1,20 +0,0 @@ -scenario: "s2" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/panoptic_val2017.json" -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" -panoptic_gt_folder: "/data/coco2017/annotations/panoptic_val2017" -gpu_id: 0 -cuda: true -test_batch_size: 1 -num_workers: 4 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py" -task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth" diff --git a/configs/eval/ctaic_s3.yaml b/configs/eval/ctaic_s3.yaml deleted file mode 100644 index ef0e1cda2c712d83d0c6752bff3cd302063dc18b..0000000000000000000000000000000000000000 --- a/configs/eval/ctaic_s3.yaml +++ /dev/null @@ -1,19 +0,0 @@ -scenario: "s3" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/coco_wholebody_val_v1.0.json" -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 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py" -task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth" diff --git a/configs/eval/taic_detection.yaml b/configs/eval/taic_detection.yaml deleted file mode 100644 index c3d16f479d63a8addee41d5bde421e2a584a7798..0000000000000000000000000000000000000000 --- a/configs/eval/taic_detection.yaml +++ /dev/null @@ -1,22 +0,0 @@ -# TAIC eval: codec test + optional --with-metrics -task: "detection" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/instances_val2017.json" -quality_level: 1 -lmbda: 0.0035 -checkpoint: "./logs/taic_detection/1/checkpoint_0.1334.pth.tar" -base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" -# null = original COCO resolution; 256 = force 256x256 (aspect stretched) -eval_size: null -# official task network (required for --with-metrics) -gpu_id: 0 -cuda: true -test_batch_size: 1 -num_workers: 4 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/detection/model_mmdet3.pth" diff --git a/configs/eval/taic_instance.yaml b/configs/eval/taic_instance.yaml deleted file mode 100644 index 6dbe1e9e148bce825603d94c43f5b8acca8387f8..0000000000000000000000000000000000000000 --- a/configs/eval/taic_instance.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# TAIC eval: codec test + optional --with-metrics -task: "instance" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/instances_val2017.json" -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" -# official task network (required for --with-metrics) -gpu_id: 0 -cuda: true -test_batch_size: 1 -num_workers: 4 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth" diff --git a/configs/eval/taic_panoptic.yaml b/configs/eval/taic_panoptic.yaml deleted file mode 100644 index 361f5a9715a6d21ae37efb82686d82e563b5b696..0000000000000000000000000000000000000000 --- a/configs/eval/taic_panoptic.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# TAIC eval: codec test + optional --with-metrics -task: "panoptic" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/panoptic_val2017.json" -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" -# official task network (required for --with-metrics) -gpu_id: 0 -cuda: true -test_batch_size: 1 -num_workers: 4 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py" -task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth" diff --git a/configs/eval/taic_pose.yaml b/configs/eval/taic_pose.yaml deleted file mode 100644 index f265af8f1afa90183ba91c60105202e2a61b32ec..0000000000000000000000000000000000000000 --- a/configs/eval/taic_pose.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# TAIC eval: codec test + optional --with-metrics -task: "pose" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/coco_wholebody_val_v1.0.json" -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" -# official task network (required for --with-metrics) -gpu_id: 0 -cuda: true -test_batch_size: 1 -num_workers: 4 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py" -task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth" diff --git a/configs/eval/taic_semantic.yaml b/configs/eval/taic_semantic.yaml deleted file mode 100644 index 8d704c89a7b18f5fd84b5b027128b4a69581e02a..0000000000000000000000000000000000000000 --- a/configs/eval/taic_semantic.yaml +++ /dev/null @@ -1,20 +0,0 @@ -# TAIC eval: codec test + optional --with-metrics -task: "semantic" -dataset_path: "/data/Dataset/coco2017" -split: "val2017" -ann_file: "annotations/panoptic_val2017.json" -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" -# official task network (required for --with-metrics) -gpu_id: 0 -cuda: true -test_batch_size: 1 -num_workers: 4 -pretrained_backbone: true - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py" -task_checkpoint: "./checkpoints/task_networks/semantic/upernet_swin_base_patch4_window7_512x512_160k_ade20k_pretrain_224x224_22K_20210526_211650-762e2178.pth" diff --git a/configs/taic/detection.yaml b/configs/taic/detection.yaml deleted file mode 100644 index 5783023f7e18786f06586877fb13ea0a5bc6a983..0000000000000000000000000000000000000000 --- a/configs/taic/detection.yaml +++ /dev/null @@ -1,29 +0,0 @@ -# FlexICM TAIC - Object Detection (Cascade Mask R-CNN + Swin-B) -root: "logs" -exp_name: "taic_detection" -task: "detection" -dataset_path: "/data/Dataset/coco2017" -base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar" # quality matches lmbda mapping -checkpoint: "/data/Dataset/FlexICM/logs/taic_detection/1/checkpoint_best_loss.pth.tar" -freeze_base_codec: false # false = fine-tune the entire TIC + SFMA + Task Connector -epochs: 999 -learning_rate: 1.0e-4 -gpu_id: 0 -quality_level: 1 -# paper lambda set: {0.0035, 0.0067, 0.0130, 0.0250, 0.2050} -lmbda: 0.2050 -num_workers: 8 -batch_size: 8 -test_batch_size: 1 -patch_size: 512 -cuda: true -save: true -seed: 42 -pretrained_backbone: true -out_channels: 128 -align_mode: "fpn" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/detection/model_mmdet3.pth" diff --git a/configs/taic/instance.yaml b/configs/taic/instance.yaml deleted file mode 100644 index 6484f3f3eb35d9004fc795441659abe4a86a726f..0000000000000000000000000000000000000000 --- a/configs/taic/instance.yaml +++ /dev/null @@ -1,27 +0,0 @@ -# FlexICM TAIC - Instance Segmentation (Cascade Mask R-CNN + Swin-B) -root: "logs" -exp_name: "taic_instance" -task: "instance" -dataset_path: "/data/Dataset/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 # Cascade Mask R-CNN + Swin-B F1 -align_mode: "fpn" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/instance/model_mmdet3.pth" diff --git a/configs/taic/panoptic.yaml b/configs/taic/panoptic.yaml deleted file mode 100644 index 0bef5b7c8b1b06509db20f278202da06e9b458cf..0000000000000000000000000000000000000000 --- a/configs/taic/panoptic.yaml +++ /dev/null @@ -1,26 +0,0 @@ -root: "logs" -exp_name: "taic_panoptic" -task: "panoptic" -dataset_path: "/data/Dataset/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" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py" -task_checkpoint: "./checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth" diff --git a/configs/taic/pose.yaml b/configs/taic/pose.yaml deleted file mode 100644 index f174b74f2fcd37979075a03b4e31586cb21bd457..0000000000000000000000000000000000000000 --- a/configs/taic/pose.yaml +++ /dev/null @@ -1,26 +0,0 @@ -root: "logs" -exp_name: "taic_pose" -task: "pose" -dataset_path: "/data/Dataset/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" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py" -task_checkpoint: "./checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth" diff --git a/configs/taic/semantic.yaml b/configs/taic/semantic.yaml deleted file mode 100644 index a394542c99bf3501e26db58444e5ccf30779c22c..0000000000000000000000000000000000000000 --- a/configs/taic/semantic.yaml +++ /dev/null @@ -1,26 +0,0 @@ -root: "logs" -exp_name: "taic_semantic" -task: "semantic" -dataset_path: "/data/Dataset/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" - -# Official task-network teacher (backbone/neck) for feature alignment D -use_official_teacher: true -task_config: "./checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py" -task_checkpoint: "./checkpoints/task_networks/semantic/upernet_swin_base_patch4_window7_512x512_160k_ade20k_pretrain_224x224_22K_20210526_211650-762e2178.pth" diff --git a/configs/task_networks/README.md b/configs/task_networks/README.md deleted file mode 100644 index 4407ad9f5c3ebf55e992f2d4a0ea59cfdf9df346..0000000000000000000000000000000000000000 --- a/configs/task_networks/README.md +++ /dev/null @@ -1,44 +0,0 @@ -# Task-network configs (for metric evaluation) - -Official detection / instance weights come from -[Swin-Transformer-Object-Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection). - -Detection and instance segmentation share the same **Cascade Mask R-CNN + Swin-B** -checkpoint; metrics differ by head output (`bbox` vs `mask`). - -| Task | Model | Official source | -|------|-------|-----------------| -| detection | **Cascade Mask R-CNN + Swin-B** | [config](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection/blob/master/configs/swin/cascade_mask_rcnn_swin_base_patch4_window7_mstrain_480-800_giou_4conv1f_adamw_3x_coco.py) / [ckpt](https://github.com/SwinTransformer/storage/releases/download/v1.0.2/cascade_mask_rcnn_swin_base_patch4_window7.pth) | -| instance | **Cascade Mask R-CNN + Swin-B** (same) | same config / checkpoint as detection | -| semantic | UPerNet + Swin-B | MMSegmentation UPerNet Swin-B | -| panoptic | MaskFormer + Swin-B | MMDetection MaskFormer Swin-B | -| pose | HigherHRNet-W32 | MMPose HigherHRNet COCO-WholeBody (HRNet backbone) | - -### Download detection / instance checkpoints - -```bash -mkdir -p checkpoints/task_networks/detection checkpoints/task_networks/instance - -# Cascade Mask R-CNN + Swin-B (shared by detection bbox + instance mask) -CKPT_URL=https://github.com/SwinTransformer/storage/releases/download/v1.0.2/cascade_mask_rcnn_swin_base_patch4_window7.pth -curl -L -o checkpoints/task_networks/detection/model.pth "$CKPT_URL" -cp checkpoints/task_networks/detection/model.pth checkpoints/task_networks/instance/model.pth -rm -f checkpoints/task_networks/detection/PLACEHOLDER checkpoints/task_networks/instance/PLACEHOLDER -``` - -Symlink or copy the official config over the stub: - -```bash -ln -sf /path/to/Swin-Transformer-Object-Detection/configs/swin/cascade_mask_rcnn_swin_base_patch4_window7_mstrain_480-800_giou_4conv1f_adamw_3x_coco.py \ - configs/task_networks/cascade_mask_rcnn_swin_base_coco.py -``` - -Eval YAML fields: - -```yaml -# detection (mAP-bbox) and instance (mAP-mask) share the same config/weights -task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py" -task_checkpoint: "./checkpoints/task_networks/detection/model.pth" # or .../instance/model.pth -``` - -**Channel note:** Swin-B F1 has 128 channels; both detection and instance TAIC use `out_channels: 128`. diff --git a/configs/task_networks/cascade_mask_rcnn_swin_base_coco.py b/configs/task_networks/cascade_mask_rcnn_swin_base_coco.py deleted file mode 100644 index 4f9c1e2ef2bd5e74ebbec3a1ed7d52935c2e4623..0000000000000000000000000000000000000000 --- a/configs/task_networks/cascade_mask_rcnn_swin_base_coco.py +++ /dev/null @@ -1,251 +0,0 @@ -# Cascade Mask R-CNN + Swin-B for FlexICM metric eval (mmdet 3.x). -# Matches the official Swin-Transformer-Object-Detection paper checkpoint: -# cascade_mask_rcnn_swin_base_patch4_window7.pth -# Use the converted weights: -# checkpoints/task_networks/detection/model_mmdet3.pth -# (produced by scripts/convert_swin_det_ckpt_to_mmdet3.py) - -_base_ = [] # self-contained for FlexICM - -model = dict( - type='CascadeRCNN', - data_preprocessor=dict( - type='DetDataPreprocessor', - mean=[123.675, 116.28, 103.53], - std=[58.395, 57.12, 57.375], - bgr_to_rgb=True, - pad_mask=True, - pad_size_divisor=32), - backbone=dict( - type='SwinTransformer', - embed_dims=128, - depths=[2, 2, 18, 2], - num_heads=[4, 8, 16, 32], - window_size=7, - mlp_ratio=4, - qkv_bias=True, - qk_scale=None, - drop_rate=0.0, - attn_drop_rate=0.0, - drop_path_rate=0.3, - patch_norm=True, - out_indices=(0, 1, 2, 3), - with_cp=False, - convert_weights=False, - init_cfg=None), - neck=dict( - type='FPN', - in_channels=[128, 256, 512, 1024], - out_channels=256, - num_outs=5), - rpn_head=dict( - type='RPNHead', - in_channels=256, - feat_channels=256, - anchor_generator=dict( - type='AnchorGenerator', - scales=[8], - ratios=[0.5, 1.0, 2.0], - strides=[4, 8, 16, 32, 64]), - bbox_coder=dict( - type='DeltaXYWHBBoxCoder', - target_means=[.0, .0, .0, .0], - target_stds=[1.0, 1.0, 1.0, 1.0]), - loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0), - loss_bbox=dict(type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=1.0)), - roi_head=dict( - type='CascadeRoIHead', - num_stages=3, - stage_loss_weights=[1, 0.5, 0.25], - bbox_roi_extractor=dict( - type='SingleRoIExtractor', - roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0), - out_channels=256, - featmap_strides=[4, 8, 16, 32]), - bbox_head=[ - dict( - type='Shared4Conv1FCBBoxHead', - in_channels=256, - conv_out_channels=256, - fc_out_channels=1024, - roi_feat_size=7, - num_classes=80, - bbox_coder=dict( - type='DeltaXYWHBBoxCoder', - target_means=[0.0, 0.0, 0.0, 0.0], - target_stds=[0.1, 0.1, 0.2, 0.2]), - reg_class_agnostic=False, - reg_decoded_bbox=True, - norm_cfg=dict(type='BN', requires_grad=True), - loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0), - loss_bbox=dict(type='GIoULoss', loss_weight=10.0)), - dict( - type='Shared4Conv1FCBBoxHead', - in_channels=256, - conv_out_channels=256, - fc_out_channels=1024, - roi_feat_size=7, - num_classes=80, - bbox_coder=dict( - type='DeltaXYWHBBoxCoder', - target_means=[0.0, 0.0, 0.0, 0.0], - target_stds=[0.05, 0.05, 0.1, 0.1]), - reg_class_agnostic=False, - reg_decoded_bbox=True, - norm_cfg=dict(type='BN', requires_grad=True), - loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0), - loss_bbox=dict(type='GIoULoss', loss_weight=10.0)), - dict( - type='Shared4Conv1FCBBoxHead', - in_channels=256, - conv_out_channels=256, - fc_out_channels=1024, - roi_feat_size=7, - num_classes=80, - bbox_coder=dict( - type='DeltaXYWHBBoxCoder', - target_means=[0.0, 0.0, 0.0, 0.0], - target_stds=[0.033, 0.033, 0.067, 0.067]), - reg_class_agnostic=False, - reg_decoded_bbox=True, - norm_cfg=dict(type='BN', requires_grad=True), - loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0), - loss_bbox=dict(type='GIoULoss', loss_weight=10.0)), - ], - mask_roi_extractor=dict( - type='SingleRoIExtractor', - roi_layer=dict(type='RoIAlign', output_size=14, sampling_ratio=0), - out_channels=256, - featmap_strides=[4, 8, 16, 32]), - mask_head=[ - dict( - type='FCNMaskHead', - num_convs=4, - in_channels=256, - conv_out_channels=256, - num_classes=80, - loss_mask=dict(type='CrossEntropyLoss', use_mask=True, loss_weight=1.0)), - dict( - type='FCNMaskHead', - num_convs=4, - in_channels=256, - conv_out_channels=256, - num_classes=80, - loss_mask=dict(type='CrossEntropyLoss', use_mask=True, loss_weight=1.0)), - dict( - type='FCNMaskHead', - num_convs=4, - in_channels=256, - conv_out_channels=256, - num_classes=80, - loss_mask=dict(type='CrossEntropyLoss', use_mask=True, loss_weight=1.0)), - ]), - train_cfg=dict( - rpn=dict( - assigner=dict( - type='MaxIoUAssigner', - pos_iou_thr=0.7, - neg_iou_thr=0.3, - min_pos_iou=0.3, - match_low_quality=True, - ignore_iof_thr=-1), - sampler=dict( - type='RandomSampler', - num=256, - pos_fraction=0.5, - neg_pos_ub=-1, - add_gt_as_proposals=False), - allowed_border=0, - pos_weight=-1, - debug=False), - rpn_proposal=dict( - nms_pre=2000, - max_per_img=2000, - nms=dict(type='nms', iou_threshold=0.7), - min_bbox_size=0), - rcnn=[ - dict( - assigner=dict( - type='MaxIoUAssigner', - pos_iou_thr=0.5, - neg_iou_thr=0.5, - min_pos_iou=0.5, - match_low_quality=False, - ignore_iof_thr=-1), - sampler=dict( - type='RandomSampler', - num=512, - pos_fraction=0.25, - neg_pos_ub=-1, - add_gt_as_proposals=True), - mask_size=28, - pos_weight=-1, - debug=False), - dict( - assigner=dict( - type='MaxIoUAssigner', - pos_iou_thr=0.6, - neg_iou_thr=0.6, - min_pos_iou=0.6, - match_low_quality=False, - ignore_iof_thr=-1), - sampler=dict( - type='RandomSampler', - num=512, - pos_fraction=0.25, - neg_pos_ub=-1, - add_gt_as_proposals=True), - mask_size=28, - pos_weight=-1, - debug=False), - dict( - assigner=dict( - type='MaxIoUAssigner', - pos_iou_thr=0.7, - neg_iou_thr=0.7, - min_pos_iou=0.7, - match_low_quality=False, - ignore_iof_thr=-1), - sampler=dict( - type='RandomSampler', - num=512, - pos_fraction=0.25, - neg_pos_ub=-1, - add_gt_as_proposals=True), - mask_size=28, - pos_weight=-1, - debug=False), - ]), - test_cfg=dict( - rpn=dict( - nms_pre=1000, - max_per_img=1000, - nms=dict(type='nms', iou_threshold=0.7), - min_bbox_size=0), - rcnn=dict( - score_thr=0.05, - nms=dict(type='nms', iou_threshold=0.5), - max_per_img=100, - mask_thr_binary=0.5))) - -# Minimal dataset meta required by mmdet.apis.init_detector -test_dataloader = dict( - dataset=dict( - type='CocoDataset', - ann_file='annotations/instances_val2017.json', - data_prefix=dict(img='val2017/'), - metainfo=dict( - classes=('person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', - 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', - 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', - 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', - 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', - 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', - 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', - 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', - 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', - 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', - 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', - 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', - 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', - 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush')))) diff --git a/configs/task_networks/higherhrnet_w32_coco_wholebody.py b/configs/task_networks/higherhrnet_w32_coco_wholebody.py deleted file mode 100644 index 6bbdceec02b1df86e878edd069d75617daf16c95..0000000000000000000000000000000000000000 --- a/configs/task_networks/higherhrnet_w32_coco_wholebody.py +++ /dev/null @@ -1,7 +0,0 @@ -# STUB: replace this file with a real OpenMMLab config (or symlink to your mmdet/mmseg/mmpose config). -# Expected model family: higherhrnet_w32_coco_wholebody -# See configs/task_networks/README.md -raise RuntimeError( - "Replace configs/task_networks/higherhrnet_w32_coco_wholebody.py with a real OpenMMLab config " - "(copy/symlink from mmdet/mmseg/mmpose)." -) diff --git a/configs/task_networks/maskformer_swin-b_coco.py b/configs/task_networks/maskformer_swin-b_coco.py deleted file mode 100644 index 010438aec4bf7a979795865576d552fd4a4b37cc..0000000000000000000000000000000000000000 --- a/configs/task_networks/maskformer_swin-b_coco.py +++ /dev/null @@ -1,7 +0,0 @@ -# STUB: replace this file with a real OpenMMLab config (or symlink to your mmdet/mmseg/mmpose config). -# Expected model family: maskformer_swin-b_coco -# See configs/task_networks/README.md -raise RuntimeError( - "Replace configs/task_networks/maskformer_swin-b_coco.py with a real OpenMMLab config " - "(copy/symlink from mmdet/mmseg/mmpose)." -) diff --git a/configs/task_networks/upernet_swin-b_coco.py b/configs/task_networks/upernet_swin-b_coco.py deleted file mode 100644 index 459410f4891b16afe190ac058ad7db10b91858cb..0000000000000000000000000000000000000000 --- a/configs/task_networks/upernet_swin-b_coco.py +++ /dev/null @@ -1,7 +0,0 @@ -# STUB: replace this file with a real OpenMMLab config (or symlink to your mmdet/mmseg/mmpose config). -# Expected model family: upernet_swin-b_coco -# See configs/task_networks/README.md -raise RuntimeError( - "Replace configs/task_networks/upernet_swin-b_coco.py with a real OpenMMLab config " - "(copy/symlink from mmdet/mmseg/mmpose)." -) diff --git a/flexicm/__init__.py b/flexicm/__init__.py deleted file mode 100644 index 3d528b8eb454360fb0b61d8ca367c4fe477d8154..0000000000000000000000000000000000000000 --- a/flexicm/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Package init for FlexICM.""" - -__version__ = "1.0.0" diff --git a/flexicm/data/__init__.py b/flexicm/data/__init__.py deleted file mode 100644 index d9c47dec18113c2650b7feaf1d891a976c88d43b..0000000000000000000000000000000000000000 --- a/flexicm/data/__init__.py +++ /dev/null @@ -1,21 +0,0 @@ -from .datasets import ( - COCOImageDataset, - COCOWholeBodyImageDataset, - ImageFolderDataset, - build_test_transform, - build_train_transform, - collate_keep, -) -from .coco_eval import COCOEvalDataset, TASK_ANN_FILES, coco_eval_collate - -__all__ = [ - "COCOImageDataset", - "COCOWholeBodyImageDataset", - "ImageFolderDataset", - "COCOEvalDataset", - "TASK_ANN_FILES", - "build_test_transform", - "build_train_transform", - "collate_keep", - "coco_eval_collate", -] diff --git a/flexicm/data/coco_eval.py b/flexicm/data/coco_eval.py deleted file mode 100644 index 098c2570ee921c3d717ce9dc24bbc5709f863bea..0000000000000000000000000000000000000000 --- a/flexicm/data/coco_eval.py +++ /dev/null @@ -1,80 +0,0 @@ -"""COCO-style evaluation datasets that return image + annotation paths/ids.""" - -from __future__ import annotations - -import json -import os -from typing import Any, Dict, List, Optional - -import torch -from PIL import Image -from torch.utils.data import Dataset -from torchvision import transforms - - -class COCOEvalDataset(Dataset): - """COCO val images with annotation ids for metric evaluation. - - Returns a dict: - image: FloatTensor CxHxW in [0,1] - image_id: int - file_name: str - height, width: int - path: str - """ - - def __init__( - self, - coco_root: str, - ann_file: str, - image_prefix: str = "val2017", - transform=None, - ): - self.coco_root = coco_root - self.image_dir = os.path.join(coco_root, image_prefix) - self.ann_file = ann_file if os.path.isabs(ann_file) else os.path.join(coco_root, ann_file) - self.transform = transform or transforms.ToTensor() - - with open(self.ann_file) as f: - coco = json.load(f) - self.images: List[Dict[str, Any]] = sorted(coco["images"], key=lambda x: x["id"]) - self.categories = coco.get("categories", []) - - def __len__(self): - return len(self.images) - - def __getitem__(self, index: int) -> Dict[str, Any]: - info = self.images[index] - path = os.path.join(self.image_dir, info["file_name"]) - img = Image.open(path).convert("RGB") - tensor = self.transform(img) - return { - "image": tensor, - "image_id": int(info["id"]), - "file_name": info["file_name"], - "height": int(info["height"]), - "width": int(info["width"]), - "path": path, - } - - -def coco_eval_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]: - """Collate that keeps variable-size images as a list (batch_size usually 1).""" - return { - "images": [b["image"] for b in batch], - "image_ids": [b["image_id"] for b in batch], - "file_names": [b["file_name"] for b in batch], - "heights": [b["height"] for b in batch], - "widths": [b["width"] for b in batch], - "paths": [b["path"] for b in batch], - } - - -# Default annotation files relative to coco_root -TASK_ANN_FILES = { - "detection": "annotations/instances_val2017.json", - "instance": "annotations/instances_val2017.json", - "semantic": "annotations/panoptic_val2017.json", # or stuff; override in config - "panoptic": "annotations/panoptic_val2017.json", - "pose": "annotations/coco_wholebody_val_v1.0.json", -} diff --git a/flexicm/data/datasets.py b/flexicm/data/datasets.py deleted file mode 100644 index 0b34d80071a75661e3eec03f13f6b6e0ebc539a4..0000000000000000000000000000000000000000 --- a/flexicm/data/datasets.py +++ /dev/null @@ -1,93 +0,0 @@ -"""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(eval_size: Optional[int] = None) -> Callable: - """Eval preprocess. If ``eval_size`` is set (e.g. 256), force HxW = size×size.""" - if eval_size is None: - return transforms.ToTensor() - return transforms.Compose( - [ - transforms.Resize((int(eval_size), int(eval_size))), - 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 deleted file mode 100644 index 5f078707434ec3670e8a1551aac4ea1c42c16df4..0000000000000000000000000000000000000000 --- a/flexicm/layers/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .layers import * -from .gdn import GDN diff --git a/flexicm/layers/gdn.py b/flexicm/layers/gdn.py deleted file mode 100644 index 099b987d561034e4ec35184c44e81485f0ccd9c4..0000000000000000000000000000000000000000 --- a/flexicm/layers/gdn.py +++ /dev/null @@ -1,121 +0,0 @@ -# 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 deleted file mode 100644 index 73712292427893af0c3873f3a06bba8711b25e51..0000000000000000000000000000000000000000 --- a/flexicm/layers/layers.py +++ /dev/null @@ -1,769 +0,0 @@ -# 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 deleted file mode 100644 index 47768a532db1c0f7b1c6d36877f1708832216d2a..0000000000000000000000000000000000000000 --- a/flexicm/models/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -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 deleted file mode 100644 index 9000811c33a65d907efff6cfa8f7cdd96568206e..0000000000000000000000000000000000000000 --- a/flexicm/models/conditional.py +++ /dev/null @@ -1,140 +0,0 @@ -"""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 deleted file mode 100644 index 3abd52277418291920775425dc81cb2e005fb3ad..0000000000000000000000000000000000000000 --- a/flexicm/models/cross_attention.py +++ /dev/null @@ -1,118 +0,0 @@ -"""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 deleted file mode 100644 index 9f2a6874e19f9375c7a73db6cbe109b3bf1556a0..0000000000000000000000000000000000000000 --- a/flexicm/models/ctaic.py +++ /dev/null @@ -1,100 +0,0 @@ -"""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 deleted file mode 100644 index b6c17027218d5c385e31b064c0dd4ab3e2713549..0000000000000000000000000000000000000000 --- a/flexicm/models/sfma.py +++ /dev/null @@ -1,40 +0,0 @@ -"""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 deleted file mode 100644 index d7161ca1538a2d2b0653682a4f94582535044003..0000000000000000000000000000000000000000 --- a/flexicm/models/taic.py +++ /dev/null @@ -1,395 +0,0 @@ -"""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, - ) - return 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 deleted file mode 100644 index 0528deec2e88e004ed094fd127639f7d383d6096..0000000000000000000000000000000000000000 --- a/flexicm/models/task_connector.py +++ /dev/null @@ -1,36 +0,0 @@ -"""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 deleted file mode 100644 index ec18be96e3981a4f9b3fbf308bc79c2c47075405..0000000000000000000000000000000000000000 --- a/flexicm/tasks/__init__.py +++ /dev/null @@ -1,310 +0,0 @@ -"""Task-specific frozen teachers for FlexICM feature alignment. - -Five tasks (paper Sec.III.A / IV.A): - 1. Object detection - Cascade Mask R-CNN + Swin-B (official Swin det zoo; mAP-bbox) - 2. Instance segmentation - Cascade Mask R-CNN + Swin-B (same zoo; mAP-mask) - 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 -from flexicm.tasks.official_teachers import OfficialHRNetTeacher, OfficialSwinTeacher - -# Default official task-network assets (used when use_official_teacher=True) -DEFAULT_TEACHER_ASSETS = { - "detection": { - "task_config": "configs/task_networks/cascade_mask_rcnn_swin_base_coco.py", - "task_checkpoint": "checkpoints/task_networks/detection/model_mmdet3.pth", - "framework": "mmdet", - "align_mode": "fpn", - }, - "instance": { - "task_config": "configs/task_networks/cascade_mask_rcnn_swin_base_coco.py", - "task_checkpoint": "checkpoints/task_networks/instance/model_mmdet3.pth", - "framework": "mmdet", - "align_mode": "fpn", - }, - "semantic": { - "task_config": "checkpoints/task_networks/semantic/swin-base-patch4-window7-in22k-pre_upernet_8xb2-160k_ade20k-512x512.py", - "task_checkpoint": "checkpoints/task_networks/semantic/upernet_swin_base_patch4_window7_512x512_160k_ade20k_pretrain_224x224_22K_20210526_211650-762e2178.pth", - "framework": "mmseg", - "align_mode": "fpn", - }, - "panoptic": { - "task_config": "checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic.py", - "task_checkpoint": "checkpoints/task_networks/panoptic/mask2former_swin-b-p4-w12-384-in21k_8xb2-lsj-50e_coco-panoptic_20220329_230021-05ec7315.pth", - "framework": "mmdet", - "align_mode": "stages", - }, - "pose": { - "task_config": "checkpoints/task_networks/pose/ae_hrnet-w32_8xb24-300e_coco-512x512.py", - "task_checkpoint": "checkpoints/task_networks/pose/hrnet_w32_coco_512x512-bcb8c247_20200816.pth", - "framework": "mmpose", - "align_mode": "stages", - }, -} - - -class DetectionTeacher(nn.Module): - """Detection / instance teacher for FPN feature alignment. - - Both tasks use Cascade Mask R-CNN + Swin-B (F1 = 128-d). - """ - - align_mode = "fpn" - out_channels = 128 - - def __init__(self, pretrained_backbone: bool = True, task: str = "detection"): - super().__init__() - self.task = task - self.backbone = SwinStageTeacher( - pretrained=pretrained_backbone, - use_fpn=True, - swin_variant="base", - ) - 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)) - use_official = kwargs.pop("use_official_teacher", True) - task_config = kwargs.pop("task_config", None) - task_checkpoint = kwargs.pop("task_checkpoint", None) - device = kwargs.pop("device", "cpu") - - # Normalize aliases - alias = { - "object_detection": "detection", - "det": "detection", - "instance_seg": "instance", - "instance_segmentation": "instance", - "semantic_seg": "semantic", - "semantic_segmentation": "semantic", - "panoptic_seg": "panoptic", - "panoptic_segmentation": "panoptic", - "pose_estimation": "pose", - } - task = alias.get(task, task) - - if use_official: - assets = DEFAULT_TEACHER_ASSETS.get(task) - if assets is None: - raise ValueError(f"Unknown task for official teacher: {task}") - cfg = task_config or assets["task_config"] - ckpt = task_checkpoint or assets["task_checkpoint"] - align_mode = kwargs.pop("align_mode", assets["align_mode"]) - framework = assets["framework"] - if framework in ("mmdet", "mmseg"): - teacher = OfficialSwinTeacher( - config_path=cfg, - checkpoint_path=ckpt, - align_mode=align_mode, - framework=framework, - device=device, - ) - freeze_module(teacher) - return teacher - if framework == "mmpose": - width = kwargs.pop("width", 32) - teacher = OfficialHRNetTeacher( - config_path=cfg, - checkpoint_path=ckpt, - device=device, - width=width, - ) - freeze_module(teacher) - return teacher - - # Legacy timm / in-repo teachers (ImageNet Swin / stem HRNet) - if task == "detection": - return DetectionTeacher(task="detection", pretrained_backbone=pretrained, **kwargs) - if task == "instance": - return DetectionTeacher(task="instance", pretrained_backbone=pretrained, **kwargs) - if task == "semantic": - return SemanticSegTeacher(pretrained_backbone=pretrained, **kwargs) - if task == "panoptic": - return PanopticSegTeacher(pretrained_backbone=pretrained, **kwargs) - if task == "pose": - 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, # Cascade Mask R-CNN + Swin-B F1 - "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 deleted file mode 100644 index e8cfb71948138aa64f074d18c83194a3c351ca46..0000000000000000000000000000000000000000 --- a/flexicm/tasks/losses.py +++ /dev/null @@ -1,74 +0,0 @@ -"""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/metric_eval.py b/flexicm/tasks/metric_eval.py deleted file mode 100644 index fb686e4a490c18f4cd0dbef930da5f588b5fd6dd..0000000000000000000000000000000000000000 --- a/flexicm/tasks/metric_eval.py +++ /dev/null @@ -1,94 +0,0 @@ -"""End-to-end codec + task-network metric evaluation loop.""" - -from __future__ import annotations - -from typing import Any, Dict, List, Optional - -import torch - -from flexicm.utils.codec_test import crop_feature_to_image, pad_for_codec -from flexicm.tasks.metric_runners import TaskMetricRunner, build_metric_runner - - -@torch.no_grad() -def run_task_metric_eval( - codec, - runner: TaskMetricRunner, - loader, - device: str, - ann_file: str, - use_condition: bool = False, - base_codec=None, - align_divisor: int = 256, - max_batches: Optional[int] = None, - finalize_kwargs: Optional[Dict[str, Any]] = None, -) -> Dict[str, float]: - """For each image: codec -> h -> truncated task net -> accumulate -> metrics.""" - codec.eval() - predictions: List[Any] = [] - - for i, batch in enumerate(loader): - if max_batches is not None and i >= max_batches: - break - - # Support both plain image batches and COCOEval collate dicts - if isinstance(batch, dict) and "images" in batch: - images_list = batch["images"] - metas = [] - for j in range(len(images_list)): - metas.append( - dict( - image_id=batch["image_ids"][j], - height=batch["heights"][j], - width=batch["widths"][j], - path=batch["paths"][j], - file_name=batch["file_names"][j], - ) - ) - else: - # Tensor batch Bx3xHxW without coco ids — skip metric (needs image_id) - raise RuntimeError( - "Task-metric eval requires COCOEvalDataset + coco_eval_collate " - "(image_id / height / width)." - ) - - for image, meta in zip(images_list, metas): - image = image.unsqueeze(0).to(device) - _, _, H, W = image.shape - ori_h = int(meta["height"]) - ori_w = int(meta["width"]) - # If the image was resized (e.g. eval_size=256), boxes must be - # mapped back with the correct scale_factor for COCO mAP. - scale_w = float(W) / float(ori_w) if ori_w > 0 else 1.0 - scale_h = float(H) / float(ori_h) if ori_h > 0 else 1.0 - x, _ = pad_for_codec(image, divisor=align_divisor, device=device) - - y_b = None - if use_condition and base_codec is not None: - y_b = base_codec(x)["y_hat"] - out = codec(x, y_b_hat=y_b, use_condition=True) - elif hasattr(codec, "forward") and use_condition is False and base_codec is None: - out = codec(x) - else: - # CTAIC without condition - if hasattr(codec, "forward"): - try: - out = codec(x, y_b_hat=None, use_condition=False) - except TypeError: - out = codec(x) - else: - out = codec(x) - - h = out["h"] - meta = dict(meta) - meta["pad_height"] = int(x.shape[-2]) - meta["pad_width"] = int(x.shape[-1]) - meta["scale_factor"] = (scale_w, scale_h) - pred = runner.predict_from_h(h, meta) - predictions.append(pred) - - if i % 20 == 0: - print(f"[metric] processed batch {i}/{len(loader)}") - - metrics = runner.finalize(predictions, ann_file, **(finalize_kwargs or {})) - return metrics diff --git a/flexicm/tasks/metric_runners.py b/flexicm/tasks/metric_runners.py deleted file mode 100644 index dea1de96368763d603d7b41783aac858069d44bc..0000000000000000000000000000000000000000 --- a/flexicm/tasks/metric_runners.py +++ /dev/null @@ -1,527 +0,0 @@ -"""Task-network metric runners: load official checkpoints and evaluate from codec feature h. - -Paper flow (Sec.III.A): - codec -> h (H/4 x W/4 x C) -> truncated task network (from Stage 2 / FPN) -> task output - then compute mAP-bbox / mAP-mask / mIoU / PQ / mAP-OKS. - -Requires optional packages: - pip install pycocotools - mim install mmdet mmsegmentation mmpose # plus mmengine mmcv -""" - -from __future__ import annotations - -import os -from abc import ABC, abstractmethod -from typing import Any, Dict, List, Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - - -class TaskMetricRunner(ABC): - """Unified interface for end-task evaluation from decoded feature h.""" - - metric_name: str = "metric" - - def __init__(self, device: str = "cuda"): - self.device = device - self.model = None - - @abstractmethod - def load(self, config_path: str, checkpoint_path: str) -> None: - ... - - @abstractmethod - def predict_from_h( - self, - h: torch.Tensor, - img_meta: Dict[str, Any], - ) -> Any: - """Run truncated task net starting from feature h.""" - ... - - @abstractmethod - def finalize(self, predictions: List[Any], ann_file: str, **kwargs) -> Dict[str, float]: - """Aggregate predictions vs GT annotations into scalar metrics.""" - ... - - -def _require_mmdet(): - try: - import mmdet # noqa: F401 - from mmdet.apis import init_detector - except ImportError as e: - raise ImportError( - "Full task-metric evaluation requires MMDetection.\n" - " pip install -U openmim && mim install mmengine mmcv mmdet" - ) from e - return init_detector - - -def _require_mmseg(): - try: - from mmseg.apis import init_model - except ImportError as e: - raise ImportError( - "Semantic segmentation metric evaluation requires MMSegmentation.\n" - " mim install mmsegmentation" - ) from e - return init_model - - -def _require_mmpose(): - try: - from mmpose.apis import init_model - except ImportError as e: - raise ImportError( - "Pose metric evaluation requires MMPose.\n" - " mim install mmpose" - ) from e - return init_model - - -def swin_feats_from_h(backbone: nn.Module, h: torch.Tensor) -> Tuple[torch.Tensor, ...]: - """Treat h as Swin F1 (stage-0 output) and run remaining stages. - - Compatible with: - - MMDet 3.x SwinTransformer (token + hw_shape API, ``stages``) - - older Swin / MMSeg variants that expose ``stages`` / ``layers`` - """ - stages = None - for name in ("stages", "layers"): - if hasattr(backbone, name): - stages = getattr(backbone, name) - break - if stages is None: - raise RuntimeError("Backbone has no stages/layers; cannot inject h as F1") - - # h is NCHW (B,C,H,W) == F1 - if h.dim() != 4: - raise ValueError(f"expected NCHW h, got shape {tuple(h.shape)}") - B, C, H, W = h.shape - outs = [h] - - # Detect mmdet3-style SwinBlockSequence: forward(x, hw_shape) - import inspect - - try: - needs_hw = "hw_shape" in inspect.signature(stages[0].forward).parameters - except (TypeError, ValueError): - needs_hw = False - - if needs_hw: - # tokens: (B, H*W, C) - x = h.flatten(2).transpose(1, 2).contiguous() - hw_shape = (H, W) - - # F1 is stage-0 block output (pre-downsample). Feed stage-0 downsample - # then run stages 1..N-1, mirroring SwinTransformer.forward. - if getattr(stages[0], "downsample", None) is not None: - x, hw_shape = stages[0].downsample(x, hw_shape) - - for i in range(1, len(stages)): - x, hw_shape, out, out_hw_shape = stages[i](x, hw_shape) - norm_name = f"norm{i}" - if hasattr(backbone, norm_name): - out = getattr(backbone, norm_name)(out) - feat_dim = out.shape[-1] - out = ( - out.view(B, out_hw_shape[0], out_hw_shape[1], feat_dim) - .permute(0, 3, 1, 2) - .contiguous() - ) - outs.append(out) - - # Optionally normalize F1 with norm0 for consistency with extract_feat - if hasattr(backbone, "norm0"): - f1 = h.flatten(2).transpose(1, 2).contiguous() - f1 = backbone.norm0(f1) - outs[0] = ( - f1.view(B, H, W, C).permute(0, 3, 1, 2).contiguous() - ) - return tuple(outs) - - # Legacy path (timm / older mmdet): stage modules accept NCHW / plain tensors - x = h - for i in range(1, len(stages)): - x = stages[i](x) - if isinstance(x, (tuple, list)): - x = x[0] - if x.dim() == 4 and x.shape[1] < x.shape[-1] and x.shape[-1] in ( - 96, 128, 192, 256, 384, 512, 768, 1024 - ): - x = x.permute(0, 3, 1, 2).contiguous() - outs.append(x) - - if hasattr(backbone, "num_features") or hasattr(backbone, "out_indices"): - norm_outs = [] - for i, out in enumerate(outs): - norm_name = f"norm{i}" - if hasattr(backbone, norm_name): - nchw = out - norm = getattr(backbone, norm_name) - try: - y = nchw.permute(0, 2, 3, 1) - y = norm(y) - nchw = y.permute(0, 3, 1, 2).contiguous() - except Exception: - nchw = out - norm_outs.append(nchw) - else: - norm_outs.append(out) - return tuple(norm_outs) - return tuple(outs) - - -class DetectionMetricRunner(TaskMetricRunner): - """Cascade Mask R-CNN + Swin-B (official zoo) — mAP-bbox / mAP-mask.""" - - def __init__(self, device: str = "cuda", with_mask: bool = False): - super().__init__(device) - self.with_mask = with_mask - self.metric_name = "mAP-mask" if with_mask else "mAP-bbox" - self._results: List[Dict] = [] - - def load(self, config_path: str, checkpoint_path: str) -> None: - init_detector = _require_mmdet() - self.model = init_detector(config_path, checkpoint_path, device=self.device) - self.model.eval() - - @torch.no_grad() - def predict_from_h(self, h: torch.Tensor, img_meta: Dict[str, Any]) -> Dict[str, Any]: - assert self.model is not None - # h: 1xCx(H/4)x(W/4) — preferably on the padded codec grid - backbone = self.model.backbone - feats = swin_feats_from_h(backbone, h) - if hasattr(self.model, "neck") and self.model.neck is not None: - feats = self.model.neck(feats) - - # Build a minimal img_metas / data_samples for mmdet 3.x or 2.x - ori_h, ori_w = int(img_meta["height"]), int(img_meta["width"]) - # If h comes from a padded codec input, prefer those spatial sizes so - # FPN strides line up; boxes are rescaled back to ori_shape. - pad_h = int(img_meta.get("pad_height", h.shape[-2] * 4)) - pad_w = int(img_meta.get("pad_width", h.shape[-1] * 4)) - sf = img_meta.get("scale_factor", (1.0, 1.0)) - if isinstance(sf, (int, float)): - scale_factor = (float(sf), float(sf)) - else: - scale_factor = (float(sf[0]), float(sf[1])) - try: - # MMDet 3.x style - from mmdet.structures import DetDataSample - - data_sample = DetDataSample() - data_sample.set_metainfo( - dict( - img_shape=(pad_h, pad_w), - ori_shape=(ori_h, ori_w), - pad_shape=(pad_h, pad_w), - scale_factor=scale_factor, - img_id=img_meta.get("image_id"), - ) - ) - rpn_results_list = self.model.rpn_head.predict(feats, [data_sample], rescale=False) - results_list = self.model.roi_head.predict( - feats, rpn_results_list, [data_sample], rescale=True - ) - pred = results_list[0] - # mmdet 3.x roi_head.predict may return DetDataSample or InstanceData - inst = pred.pred_instances if hasattr(pred, "pred_instances") else pred - out = { - "image_id": img_meta["image_id"], - "bboxes": inst.bboxes.detach().cpu(), - "scores": inst.scores.detach().cpu(), - "labels": inst.labels.detach().cpu(), - } - if self.with_mask and hasattr(inst, "masks") and inst.masks is not None: - out["masks"] = inst.masks.to_ndarray() if hasattr(inst.masks, "to_ndarray") else inst.masks.detach().cpu() - return out - except Exception: - # Fallback MMDet 2.x - img_metas = [ - dict( - img_shape=(pad_h, pad_w, 3), - ori_shape=(ori_h, ori_w, 3), - pad_shape=(pad_h, pad_w, 3), - scale_factor=scale_factor, - flip=False, - ) - ] - proposal_list = self.model.rpn_head.simple_test_rpn(feats, img_metas) - det_results = self.model.roi_head.simple_test( - feats, proposal_list, img_metas, rescale=True - ) - # det_results: list of (bboxes_per_class) or (bboxes, segm) - return {"image_id": img_meta["image_id"], "raw": det_results[0]} - - def finalize(self, predictions: List[Any], ann_file: str, **kwargs) -> Dict[str, float]: - from pycocotools.coco import COCO - from pycocotools.cocoeval import COCOeval - import numpy as np - - coco_gt = COCO(ann_file) - coco_results = [] - for pred in predictions: - if pred is None: - continue - if "raw" in pred: - # mmdet 2.x format: list[ndarray(n,5)] per class - raw = pred["raw"] - bbox_results = raw[0] if isinstance(raw, tuple) else raw - for label, bboxes in enumerate(bbox_results): - for row in bboxes: - x1, y1, x2, y2, score = row[:5] - coco_results.append( - { - "image_id": int(pred["image_id"]), - "category_id": int(coco_gt.getCatIds()[label]) - if label < len(coco_gt.getCatIds()) - else int(label + 1), - "bbox": [float(x1), float(y1), float(x2 - x1), float(y2 - y1)], - "score": float(score), - } - ) - continue - - bboxes = pred["bboxes"].numpy() - scores = pred["scores"].numpy() - labels = pred["labels"].numpy() - cat_ids = coco_gt.getCatIds() - for box, score, label in zip(bboxes, scores, labels): - x1, y1, x2, y2 = box.tolist() - cat_id = int(cat_ids[int(label)]) if int(label) < len(cat_ids) else int(label) + 1 - coco_results.append( - { - "image_id": int(pred["image_id"]), - "category_id": cat_id, - "bbox": [x1, y1, x2 - x1, y2 - y1], - "score": float(score), - } - ) - - if not coco_results: - return {self.metric_name: 0.0} - - coco_dt = coco_gt.loadRes(coco_results) - coco_eval = COCOeval(coco_gt, coco_dt, iouType="bbox") - coco_eval.evaluate() - coco_eval.accumulate() - coco_eval.summarize() - metrics = {"mAP-bbox": float(coco_eval.stats[0])} - - if self.with_mask: - # Mask eval requires segmentation results in COCO format; if unavailable, skip - try: - coco_eval_m = COCOeval(coco_gt, coco_dt, iouType="segm") - coco_eval_m.evaluate() - coco_eval_m.accumulate() - coco_eval_m.summarize() - metrics["mAP-mask"] = float(coco_eval_m.stats[0]) - except Exception as e: - metrics["mAP-mask_error"] = str(e) - return metrics - - -class SemanticMetricRunner(TaskMetricRunner): - """UPerNet (MMSeg) — metric: mIoU.""" - - metric_name = "mIoU" - - def load(self, config_path: str, checkpoint_path: str) -> None: - init_model = _require_mmseg() - self.model = init_model(config_path, checkpoint_path, device=self.device) - self.model.eval() - self._preds = [] - - @torch.no_grad() - def predict_from_h(self, h: torch.Tensor, img_meta: Dict[str, Any]) -> Dict[str, Any]: - assert self.model is not None - backbone = self.model.backbone - feats = swin_feats_from_h(backbone, h) - seg_logits = self.model.decode_head(feats) - if isinstance(seg_logits, (tuple, list)): - seg_logits = seg_logits[0] - H, W = int(img_meta["height"]), int(img_meta["width"]) - seg = F.interpolate(seg_logits, size=(H, W), mode="bilinear", align_corners=False) - pred = seg.argmax(dim=1)[0].detach().cpu().numpy() - return {"image_id": img_meta["image_id"], "seg": pred, "path": img_meta.get("path")} - - def finalize(self, predictions: List[Any], ann_file: str, **kwargs) -> Dict[str, float]: - """Compute mIoU if GT semantic maps are provided via kwargs['gt_dir'] or panoptic conversion. - - For a minimal working path, expects kwargs['gt_seg_loader'](image_id)->HxW label map. - """ - gt_loader = kwargs.get("gt_seg_loader") - if gt_loader is None: - return { - "mIoU": float("nan"), - "note": "Provide gt_seg_loader or use panoptic stuff GT to compute mIoU", - } - - import numpy as np - - num_classes = int(kwargs.get("num_classes", 133)) - intersect = np.zeros(num_classes, dtype=np.float64) - union = np.zeros(num_classes, dtype=np.float64) - for pred in predictions: - gt = gt_loader(pred["image_id"]) - pr = pred["seg"] - if gt.shape != pr.shape: - # nearest resize pred already at image size; skip mismatch - continue - for c in range(num_classes): - pb = pr == c - gb = gt == c - inter = np.logical_and(pb, gb).sum() - uni = np.logical_or(pb, gb).sum() - intersect[c] += inter - union[c] += uni - ious = intersect / np.maximum(union, 1) - valid = union > 0 - miou = float(ious[valid].mean()) if valid.any() else 0.0 - return {"mIoU": miou} - - -class PanopticMetricRunner(TaskMetricRunner): - """MaskFormer (MMDet) — metric: PQ.""" - - metric_name = "PQ" - - def load(self, config_path: str, checkpoint_path: str) -> None: - init_detector = _require_mmdet() - self.model = init_detector(config_path, checkpoint_path, device=self.device) - self.model.eval() - - @torch.no_grad() - def predict_from_h(self, h: torch.Tensor, img_meta: Dict[str, Any]) -> Dict[str, Any]: - assert self.model is not None - # MaskFormer typically uses backbone features F1..F4 directly - backbone = self.model.backbone - feats = swin_feats_from_h(backbone, h) - H, W = int(img_meta["height"]), int(img_meta["width"]) - try: - from mmdet.structures import DetDataSample - - data_sample = DetDataSample() - data_sample.set_metainfo( - dict(img_shape=(H, W), ori_shape=(H, W), pad_shape=(H, W), img_id=img_meta["image_id"]) - ) - # panoptic head path differs by version; store feats for custom head call - if hasattr(self.model, "panoptic_head"): - results = self.model.panoptic_head.predict(feats, [data_sample], rescale=True) - return {"image_id": img_meta["image_id"], "panoptic": results[0]} - if hasattr(self.model, "simple_test"): - # older API expects image tensor; not ideal for h-injection - return {"image_id": img_meta["image_id"], "feats_only": True, "error": "need panoptic_head"} - except Exception as e: - return {"image_id": img_meta["image_id"], "error": str(e)} - return {"image_id": img_meta["image_id"], "error": "unsupported MaskFormer API"} - - def finalize(self, predictions: List[Any], ann_file: str, **kwargs) -> Dict[str, float]: - # Full PQ needs panopticapi; keep a clear placeholder result if preds incomplete - try: - from panopticapi.evaluation import pq_compute - except ImportError: - return { - "PQ": float("nan"), - "note": "Install panopticapi and provide GT panoptic folder to compute PQ", - } - gt_folder = kwargs.get("gt_folder") - pred_folder = kwargs.get("pred_folder") - if not gt_folder or not pred_folder: - return {"PQ": float("nan"), "note": "Need gt_folder and pred_folder for pq_compute"} - results = pq_compute(ann_file, kwargs.get("pred_json"), gt_folder, pred_folder) - return {"PQ": float(results["All"]["pq"])} - - -class PoseMetricRunner(TaskMetricRunner): - """HigherHRNet (MMPose, original HRNet backbone) — metric: mAP-OKS.""" - - metric_name = "mAP-OKS" - - def load(self, config_path: str, checkpoint_path: str) -> None: - init_model = _require_mmpose() - self.model = init_model(config_path, checkpoint_path, device=self.device) - self.model.eval() - - @torch.no_grad() - def predict_from_h(self, h: torch.Tensor, img_meta: Dict[str, Any]) -> Dict[str, Any]: - """Inject h as early HRNet feature when possible; else return error guidance. - - HigherHRNet uses HRNet (not Swin). Codec `out_channels` should match stem width - (default 32). Full keypoint head wiring depends on mmpose version. - """ - assert self.model is not None - try: - # Best-effort: if backbone has stage transitions, set first stream feature to h - backbone = self.model.backbone if hasattr(self.model, "backbone") else self.model - # Many mmpose models expect full image; document limitation - if hasattr(self.model, "predict"): - # Without image path, we only support feature injection hooks if present - return { - "image_id": img_meta["image_id"], - "error": ( - "HigherHRNet-from-h requires a project-specific backbone hook; " - "set pose.eval_from_image=true in config to run image-based fallback " - "after optional RGB decode, or implement HRNet stem replacement." - ), - } - except Exception as e: - return {"image_id": img_meta["image_id"], "error": str(e)} - return {"image_id": img_meta["image_id"], "error": "pose from-h not hooked"} - - def finalize(self, predictions: List[Any], ann_file: str, **kwargs) -> Dict[str, float]: - # Standard COCO keypoint eval when predictions are in COCO format - valid = [p for p in predictions if p and "keypoints" in p] - if not valid: - return { - "mAP-OKS": float("nan"), - "note": "No keypoint predictions; implement HigherHRNet-from-h or provide COCO-format preds", - } - from pycocotools.coco import COCO - from pycocotools.cocoeval import COCOeval - - coco_gt = COCO(ann_file) - coco_dt = coco_gt.loadRes(valid) - ev = COCOeval(coco_gt, coco_dt, iouType="keypoints") - ev.evaluate() - ev.accumulate() - ev.summarize() - return {"mAP-OKS": float(ev.stats[0])} - - -def build_metric_runner(task: str, device: str = "cuda") -> TaskMetricRunner: - task = task.lower() - if task in ("detection", "det", "object_detection"): - return DetectionMetricRunner(device=device, with_mask=False) - if task in ("instance", "instance_seg", "instance_segmentation"): - return DetectionMetricRunner(device=device, with_mask=True) - if task in ("semantic", "semantic_seg", "semantic_segmentation"): - return SemanticMetricRunner(device=device) - if task in ("panoptic", "panoptic_seg", "panoptic_segmentation"): - return PanopticMetricRunner(device=device) - if task in ("pose", "pose_estimation"): - return PoseMetricRunner(device=device) - raise ValueError(f"Unknown task for metric runner: {task}") - - -# Suggested OpenMMLab config names (user must download matching weights) -DEFAULT_TASK_NET_CONFIGS = { - # Official Swin-Transformer-Object-Detection zoo (same Cascade Mask R-CNN + Swin-B) - "detection": "configs/task_networks/cascade_mask_rcnn_swin_base_coco.py", # mAP-bbox - "instance": "configs/task_networks/cascade_mask_rcnn_swin_base_coco.py", # mAP-mask - "semantic": "configs/task_networks/upernet_swin-b_coco.py", - "panoptic": "configs/task_networks/maskformer_swin-b_coco.py", - "pose": "configs/task_networks/higherhrnet_w32_coco_wholebody.py", -} - -DEFAULT_TASK_NET_CKPTS = { - "detection": "checkpoints/task_networks/detection/model.pth", - "instance": "checkpoints/task_networks/instance/model.pth", - "semantic": "checkpoints/task_networks/semantic/model.pth", - "panoptic": "checkpoints/task_networks/panoptic/model.pth", - "pose": "checkpoints/task_networks/pose/model.pth", -} diff --git a/flexicm/tasks/official_teachers.py b/flexicm/tasks/official_teachers.py deleted file mode 100644 index 5f1ffeb6557e93d767fa72b80a80ea3fe8c2800a..0000000000000000000000000000000000000000 --- a/flexicm/tasks/official_teachers.py +++ /dev/null @@ -1,279 +0,0 @@ -"""Teachers that load official task-network backbones (+ necks) for feature alignment. - -Training Distortion D is computed against these frozen features so that codec -outputs match the same backbone used at metric evaluation time. -""" - -from __future__ import annotations - -import os -from typing import Dict, Optional, Tuple - -import torch -import torch.nn as nn -import torch.nn.functional as F - -from flexicm.tasks.losses import freeze_module -from flexicm.tasks.metric_runners import swin_feats_from_h - - -def _resolve_path(path: Optional[str], repo_root: Optional[str] = None) -> Optional[str]: - if not path: - return None - if os.path.isabs(path): - return path - if repo_root is None: - repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) - return os.path.join(repo_root, path) - - -def _imagenet_norm_rgb01(x: torch.Tensor) -> torch.Tensor: - """Normalize RGB float images in [0, 1] with ImageNet mean/std (float space).""" - 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 _mmdet_preprocess_rgb01(x: torch.Tensor, model: nn.Module) -> torch.Tensor: - """Apply DetDataPreprocessor-equivalent norm: RGB [0,1] -> model input.""" - pp = getattr(model, "data_preprocessor", None) - if pp is None: - return _imagenet_norm_rgb01(x) - # DetDataPreprocessor stores mean/std for 0-255 inputs - mean = getattr(pp, "mean", None) - std = getattr(pp, "std", None) - if mean is None or std is None: - return _imagenet_norm_rgb01(x) - mean_t = mean.view(1, -1, 1, 1).to(device=x.device, dtype=x.dtype) - std_t = std.view(1, -1, 1, 1).to(device=x.device, dtype=x.dtype) - # x in [0,1] RGB; preprocessor mean/std are for 0-255 RGB after bgr_to_rgb - return (x * 255.0 - mean_t) / std_t - - -class OfficialSwinTeacher(nn.Module): - """Frozen Swin teacher from an official MMDet / MMSeg checkpoint. - - align_mode: - - ``fpn``: P2..P6 via backbone + FPN/neck (detection / instance / semantic) - - ``stages``: F1..F4 via backbone stages (panoptic) - """ - - def __init__( - self, - config_path: str, - checkpoint_path: str, - align_mode: str = "fpn", - framework: str = "mmdet", - device: str = "cpu", - ): - super().__init__() - assert align_mode in ("fpn", "stages") - self.align_mode = align_mode - self.framework = framework - self.config_path = _resolve_path(config_path) - self.checkpoint_path = _resolve_path(checkpoint_path) - if not self.config_path or not os.path.isfile(self.config_path): - raise FileNotFoundError(f"task_config not found: {config_path}") - if not self.checkpoint_path or not os.path.isfile(self.checkpoint_path): - raise FileNotFoundError(f"task_checkpoint not found: {checkpoint_path}") - - self.model = self._load_model(device) - freeze_module(self.model) - self.out_channels = self._infer_f1_channels() - self._aux_fpn = None - - def _load_model(self, device: str) -> nn.Module: - if self.framework == "mmdet": - from mmdet.apis import init_detector - - return init_detector(self.config_path, self.checkpoint_path, device=device) - if self.framework == "mmseg": - from mmseg.apis import init_model - - return init_model(self.config_path, self.checkpoint_path, device=device) - raise ValueError(f"Unknown framework={self.framework}") - - def _infer_f1_channels(self) -> int: - bb = self.model.backbone - if hasattr(bb, "num_features"): - return int(bb.num_features[0]) - if hasattr(bb, "embed_dims"): - return int(bb.embed_dims) - return 128 - - def _preprocess(self, images: torch.Tensor) -> torch.Tensor: - return _mmdet_preprocess_rgb01(images, self.model) - - def _backbone_stages_from_image(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]: - """Return F1..F4 NCHW from Swin backbone (with per-stage norms).""" - bb = self.model.backbone - # Prefer official forward when possible (handles abs pos / dropout) - if hasattr(bb, "forward") and not hasattr(bb, "stages"): - feats = bb(x) - return tuple(feats[:4]) - - tokens, hw_shape = bb.patch_embed(x) - if getattr(bb, "use_abs_pos_embed", False): - tokens = tokens + bb.absolute_pos_embed - tokens = bb.drop_after_pos(tokens) - outs = [] - for i, stage in enumerate(bb.stages): - tokens, hw_shape, out, out_hw = stage(tokens, hw_shape) - if i in getattr(bb, "out_indices", (0, 1, 2, 3)): - norm = getattr(bb, f"norm{i}") - out = norm(out) - feat = ( - out.view(-1, *out_hw, bb.num_features[i]) - .permute(0, 3, 1, 2) - .contiguous() - ) - outs.append(feat) - return tuple(outs) - - def _fpn_from_stages(self, stages: Tuple[torch.Tensor, ...]) -> Dict[str, torch.Tensor]: - neck = getattr(self.model, "neck", None) - if neck is None and self.framework == "mmseg": - # UPerNet has no separate FPN neck; reuse Cascade Mask R-CNN FPN - # (same Swin-B channel layout) so Eq.2 aligns P2..P6. - neck = self._get_or_build_aux_fpn(stages[0].device) - if neck is None: - p = list(stages[:4]) - while len(p) < 4: - p.append(p[-1]) - p2, p3, p4, p5 = p - p6 = F.avg_pool2d(p5, kernel_size=2, stride=2) - return {"p2": p2, "p3": p3, "p4": p4, "p5": p5, "p6": p6} - - pyramid = neck(stages) - keys = ["p2", "p3", "p4", "p5", "p6"] - out = {} - for i, feat in enumerate(pyramid): - if i < len(keys): - out[keys[i]] = feat - return out - - def _get_or_build_aux_fpn(self, device) -> nn.Module: - if hasattr(self, "_aux_fpn") and self._aux_fpn is not None: - return self._aux_fpn - from mmdet.models.necks import FPN - - fpn = FPN( - in_channels=[128, 256, 512, 1024], - out_channels=256, - num_outs=5, - ) - # Load official Cascade Mask R-CNN neck weights when available - det_ckpt = _resolve_path("checkpoints/task_networks/detection/model_mmdet3.pth") - if det_ckpt and os.path.isfile(det_ckpt): - raw = torch.load(det_ckpt, map_location="cpu") - state = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw - neck_state = { - k[len("neck.") :]: v for k, v in state.items() if k.startswith("neck.") - } - missing, unexpected = fpn.load_state_dict(neck_state, strict=False) - # missing/unexpected are fine to ignore for BN tracking buffers etc. - self._aux_fpn = freeze_module(fpn.to(device)) - return self._aux_fpn - - def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]: - x = self._preprocess(images) - if ( - self.align_mode == "fpn" - and hasattr(self.model, "extract_feat") - and getattr(self.model, "neck", None) is not None - ): - # CascadeRCNN.extract_feat = neck(backbone(x)) -> P2..P6 - feats = self.model.extract_feat(x) - keys = ["p2", "p3", "p4", "p5", "p6"] - return {keys[i]: feats[i] for i in range(min(len(keys), len(feats)))} - - if self.framework == "mmseg" and hasattr(self.model, "extract_feat"): - stages = tuple(self.model.extract_feat(x)[:4]) - else: - stages = self._backbone_stages_from_image(x) - if self.align_mode == "stages": - return {f"f{i+1}": stages[i] for i in range(min(4, len(stages)))} - return self._fpn_from_stages(stages) - - def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: - stages = swin_feats_from_h(self.model.backbone, h) - if self.align_mode == "stages": - return {f"f{i+1}": stages[i] for i in range(min(4, len(stages)))} - return self._fpn_from_stages(stages) - - -class OfficialHRNetTeacher(nn.Module): - """Frozen HRNet / HigherHRNet-style teacher from MMPose (stages F1..F4).""" - - align_mode = "stages" - out_channels = 32 - - def __init__( - self, - config_path: str, - checkpoint_path: str, - device: str = "cpu", - width: int = 32, - ): - super().__init__() - self.config_path = _resolve_path(config_path) - self.checkpoint_path = _resolve_path(checkpoint_path) - if not self.config_path or not os.path.isfile(self.config_path): - raise FileNotFoundError(f"task_config not found: {config_path}") - if not self.checkpoint_path or not os.path.isfile(self.checkpoint_path): - raise FileNotFoundError(f"task_checkpoint not found: {checkpoint_path}") - - from mmpose.apis import init_model - - self.model = init_model(self.config_path, self.checkpoint_path, device=device) - freeze_module(self.model) - self.width = width - self.out_channels = width - # projection when codec h channels != stem width - self.h_proj = freeze_module(nn.Conv2d(128, width, 1)) - - def _preprocess(self, images: torch.Tensor) -> torch.Tensor: - return _imagenet_norm_rgb01(images) - - def _stages_from_backbone(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: - bb = self.model.backbone if hasattr(self.model, "backbone") else self.model - # HRNet forward usually returns multi-resolution list / tensor - feats = bb(x) - if isinstance(feats, (list, tuple)): - # take highest-res stream as F1, then downsample proxies for F2..F4 - f1 = feats[0] if feats[0].dim() == 4 else feats[0][-1] - outs = [f1] - cur = f1 - for i in range(1, 4): - if i < len(feats) and isinstance(feats[i], torch.Tensor) and feats[i].dim() == 4: - outs.append(feats[i]) - cur = feats[i] - else: - cur = F.avg_pool2d(cur, 2, 2) - outs.append(cur) - return {f"f{i+1}": outs[i] for i in range(4)} - # single tensor: synthesize pyramid - f1 = feats - outs = [f1] - cur = f1 - for _ in range(3): - cur = F.avg_pool2d(cur, 2, 2) - outs.append(cur) - return {f"f{i+1}": outs[i] for i in range(4)} - - def gt_features(self, images: torch.Tensor) -> Dict[str, torch.Tensor]: - return self._stages_from_backbone(self._preprocess(images)) - - def pred_features(self, h: torch.Tensor) -> Dict[str, torch.Tensor]: - if h.shape[1] != self.width: - if self.h_proj.in_channels != h.shape[1]: - self.h_proj = freeze_module(nn.Conv2d(h.shape[1], self.width, 1).to(h.device)) - h = self.h_proj(h) - # Approximate remaining stages with stride-2 pools (HRNet full from-h - # needs model-specific stem hooks; pyramid still provides multi-scale D). - outs = [h] - cur = h - for _ in range(3): - cur = F.avg_pool2d(cur, 2, 2) - outs.append(cur) - return {f"f{i+1}": outs[i] for i in range(4)} diff --git a/flexicm/tasks/swin_teacher.py b/flexicm/tasks/swin_teacher.py deleted file mode 100644 index 62c1b87299303ec2060630e32b07dd12fc3c6321..0000000000000000000000000000000000000000 --- a/flexicm/tasks/swin_teacher.py +++ /dev/null @@ -1,250 +0,0 @@ -"""Swin backbone helpers shared by detection / segmentation teachers. - -Swin-B: F1 at H/4 with C=128 (Cascade Mask R-CNN / UPerNet / MaskFormer). -Optional Swin-T/S variants are supported via `swin_variant` for experiments. - -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 - -_SWIN_TIMM_NAMES = { - "base": "swin_base_patch4_window7_224", - "tiny": "swin_tiny_patch4_window7_224", - "small": "swin_small_patch4_window7_224", -} - - -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_backbone(pretrained: bool = True, swin_variant: str = "base"): - """Build Swin via timm; returns backbone with features_only stages.""" - try: - import timm - except ImportError as e: - raise ImportError("Please install timm to use Swin teachers: pip install timm") from e - - key = swin_variant.lower().replace("swin-", "").replace("swin_", "") - if key not in _SWIN_TIMM_NAMES: - raise ValueError(f"Unknown swin_variant={swin_variant!r}; expected one of {list(_SWIN_TIMM_NAMES)}") - - # Training uses 256x256 crops; validation pads to multiples of 256. - # Allow variable spatial sizes (do not hard-lock PatchEmbed to 224). - model = timm.create_model( - _SWIN_TIMM_NAMES[key], - pretrained=pretrained, - features_only=True, - out_indices=(0, 1, 2, 3), - img_size=256, - dynamic_img_size=True, - strict_img_size=False, - ) - return model - - -def build_swin_b_backbone(pretrained: bool = True): - """Backward-compatible alias for Swin-B.""" - return build_swin_backbone(pretrained=pretrained, swin_variant="base") - - -class SwinStageTeacher(nn.Module): - """ - Extract F1..F4 from a Swin backbone (base / tiny / small). - Truncated path: treat input h as F1, run remaining stages. - """ - - def __init__( - self, - pretrained: bool = True, - use_fpn: bool = True, - fpn_dim: int = 256, - swin_variant: str = "base", - ): - super().__init__() - self.swin_variant = swin_variant - self.backbone = freeze_module( - build_swin_backbone(pretrained=pretrained, swin_variant=swin_variant) - ) - # timm: Swin-B [128,256,512,1024], Swin-T [96,192,384,768] - self.feat_channels = list(self.backbone.feature_info.channels()) - self.use_fpn = use_fpn - if use_fpn: - self.fpn = freeze_module(SimpleFPN(self.feat_channels, fpn_dim=fpn_dim)) - else: - self.fpn = None - - 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 deleted file mode 100644 index 8fbd81c8ee5b87de6fed64275840eb5273515203..0000000000000000000000000000000000000000 --- a/flexicm/utils/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -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 deleted file mode 100644 index 777ca53d6d2521c4b567427550179922ae322ff2..0000000000000000000000000000000000000000 --- a/flexicm/utils/alignment.py +++ /dev/null @@ -1,88 +0,0 @@ -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/codec_test.py b/flexicm/utils/codec_test.py deleted file mode 100644 index b2caf56aa694af3b48a9958562ab3d6e9492250c..0000000000000000000000000000000000000000 --- a/flexicm/utils/codec_test.py +++ /dev/null @@ -1,227 +0,0 @@ -"""Shared helpers for codec test / eval scripts.""" - -from __future__ import annotations - -import os -from typing import Dict, Optional, Tuple - -import torch - -from flexicm.utils.alignment import Alignment -from flexicm.utils.train_utils import AverageMeter - - -def resolve_ckpt(path: str, repo_root: str, label: str = "checkpoint") -> str: - if not path: - raise FileNotFoundError(f"{label}: 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}: 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 file: {path}") - if os.path.basename(path) == "PLACEHOLDER" or path.endswith(".txt"): - raise FileNotFoundError(f"{label}: refusing placeholder file: {path}") - return path - - -def crop_feature_to_image(h: torch.Tensor, image_hw: Tuple[int, int]) -> torch.Tensor: - """Crop decoded feature h (H/4, W/4 of padded input) to original image size / 4.""" - H, W = image_hw - return h[..., : H // 4, : W // 4] - - -def pad_for_codec(images: torch.Tensor, divisor: int = 256, device=None): - align = Alignment(divisor=divisor, mode="pad", padding_mode="constant") - if device is not None: - align = align.to(device) - return align.align(images), align - - -def likelihood_bpp(likelihoods: Dict[str, torch.Tensor], num_pixels: int) -> torch.Tensor: - import math - - return sum( - (torch.log(lik).sum() / (-math.log(2) * num_pixels)) - for lik in likelihoods.values() - ) - - -def actual_bitstream_bpp(strings, num_pixels: int) -> float: - """Estimate bpp from CompressAI byte strings: [[y_bytes...], [z_bytes...]].""" - total_bits = 0 - for group in strings: - for s in group: - if isinstance(s, (bytes, bytearray)): - total_bits += len(s) * 8 - elif torch.is_tensor(s): - total_bits += int(s.numel() * s.element_size() * 8) - else: - total_bits += len(s) * 8 - return total_bits / float(num_pixels) - - -@torch.no_grad() -def test_taic_loader( - model, - teacher, - loader, - criterion, - device, - align_divisor: int = 256, - run_actual_bpp: bool = False, - max_batches: Optional[int] = None, - log_every: int = 50, -): - """Run codec test: likelihood bpp + feature distortion (+ optional real bpp).""" - model.eval() - teacher.eval() - meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion", "actual_bpp")} - - if run_actual_bpp: - model.update(force=True) - - for i, images in enumerate(loader): - if max_batches is not None and i >= max_batches: - break - images = images.to(device) - N, _, H, W = images.shape - num_pixels = N * H * W - - # Keep codec + teacher on the padded grid so Swin always sees a - # patch/window-divisible size. bpp still uses the original pixel count. - x, _ = pad_for_codec(images, divisor=align_divisor, device=device) - out = model(x) - - gt = teacher.gt_features(x) - pred = teacher.pred_features(out["h"]) - stats = criterion(out, pred, gt, num_pixels=num_pixels) - - meters["loss"].update(stats["loss"].item(), n=N) - meters["bpp"].update(stats["bpp"].item(), n=N) - meters["distortion"].update(stats["distortion"].item(), n=N) - - if run_actual_bpp: - try: - enc = model.compress(x) - dec = model.decompress( - enc["strings"], enc["shape"], x_size=(x.shape[2], x.shape[3]) - ) - abpp = actual_bitstream_bpp(enc["strings"], num_pixels) - meters["actual_bpp"].update(abpp, n=N) - # sanity: decoded h spatial size - _ = dec["h"] - except Exception as e: - if i == 0: - print(f"[warn] actual bpp / compress-decompress failed: {e}") - - if i % log_every == 0: - msg = ( - f"[{i}/{len(loader)}] bpp={meters['bpp'].avg:.4f} " - f"D={meters['distortion'].avg:.6f} loss={meters['loss'].avg:.4f}" - ) - if run_actual_bpp and meters["actual_bpp"].count > 0: - msg += f" actual_bpp={meters['actual_bpp'].avg:.4f}" - print(msg) - - result = { - "bpp": meters["bpp"].avg, - "distortion": meters["distortion"].avg, - "loss": meters["loss"].avg, - "num_batches": meters["bpp"].count, - } - if run_actual_bpp and meters["actual_bpp"].count > 0: - result["actual_bpp"] = meters["actual_bpp"].avg - return result - - -@torch.no_grad() -def test_ctaic_loader( - ext_model, - base_model, - teacher, - loader, - criterion, - device, - use_condition: bool = True, - align_divisor: int = 256, - run_actual_bpp: bool = False, - max_batches: Optional[int] = None, - log_every: int = 50, -): - """Codec test for C-TAIC; bpp is extension-layer only (paper Sec.IV.E.2).""" - ext_model.eval() - base_model.eval() - teacher.eval() - meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion", "actual_bpp")} - - if run_actual_bpp: - ext_model.update(force=True) - - for i, images in enumerate(loader): - if max_batches is not None and i >= max_batches: - break - images = images.to(device) - N, _, H, W = images.shape - num_pixels = N * H * W - - # Keep codec + teacher on the padded grid (same rationale as TAIC test). - x, _ = pad_for_codec(images, divisor=align_divisor, device=device) - y_b = None - if use_condition: - base_out = base_model(x) - y_b = base_out["y_hat"] - - out = ext_model(x, y_b_hat=y_b, use_condition=use_condition and y_b is not None) - - gt = teacher.gt_features(x) - pred = teacher.pred_features(out["h"]) - stats = criterion(out, pred, gt, num_pixels=num_pixels) - - meters["loss"].update(stats["loss"].item(), n=N) - meters["bpp"].update(stats["bpp"].item(), n=N) - meters["distortion"].update(stats["distortion"].item(), n=N) - - if run_actual_bpp: - try: - enc = ext_model.compress( - x, y_b_hat=y_b, use_condition=use_condition and y_b is not None - ) - dec = ext_model.decompress( - enc["strings"], - enc["shape"], - x_size=(x.shape[2], x.shape[3]), - y_b_hat=y_b, - use_condition=use_condition and y_b is not None, - ) - abpp = actual_bitstream_bpp(enc["strings"], num_pixels) - meters["actual_bpp"].update(abpp, n=N) - _ = dec["h"] - except Exception as e: - if i == 0: - print(f"[warn] actual bpp / compress-decompress failed: {e}") - - if i % log_every == 0: - msg = ( - f"[{i}/{len(loader)}] bpp={meters['bpp'].avg:.4f} " - f"D={meters['distortion'].avg:.6f} loss={meters['loss'].avg:.4f}" - ) - if run_actual_bpp and meters["actual_bpp"].count > 0: - msg += f" actual_bpp={meters['actual_bpp'].avg:.4f}" - print(msg) - - result = { - "bpp": meters["bpp"].avg, - "distortion": meters["distortion"].avg, - "loss": meters["loss"].avg, - "num_batches": meters["bpp"].count, - "use_condition": use_condition, - } - if run_actual_bpp and meters["actual_bpp"].count > 0: - result["actual_bpp"] = meters["actual_bpp"].avg - return result diff --git a/flexicm/utils/dataloader.py b/flexicm/utils/dataloader.py deleted file mode 100644 index 3d8c9a330f7a6a26731df508cd35912015d291eb..0000000000000000000000000000000000000000 --- a/flexicm/utils/dataloader.py +++ /dev/null @@ -1,68 +0,0 @@ -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 deleted file mode 100644 index e57f94f747c507a15f5c9e66c1280e3a74a19cf9..0000000000000000000000000000000000000000 --- a/flexicm/utils/train_utils.py +++ /dev/null @@ -1,96 +0,0 @@ -"""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 deleted file mode 100644 index fbcd4c20e01d08988a79b8cc1bb11e464632b701..0000000000000000000000000000000000000000 --- a/requirements.txt +++ /dev/null @@ -1,19 +0,0 @@ -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 -pycocotools>=2.0.6 - -# Optional — required for --with-metrics (full task-network evaluation) -# openmim -# mmengine -# mmcv -# mmdet -# mmsegmentation -# mmpose -# panopticapi # for PQ diff --git a/scripts/convert_swin_det_ckpt_to_mmdet3.py b/scripts/convert_swin_det_ckpt_to_mmdet3.py deleted file mode 100644 index 9bf06f4295a64b88401481096382b332efe239ac..0000000000000000000000000000000000000000 --- a/scripts/convert_swin_det_ckpt_to_mmdet3.py +++ /dev/null @@ -1,63 +0,0 @@ -#!/usr/bin/env python3 -"""Convert official Swin-Det Cascade Mask R-CNN ckpt -> mmdet 3.x key layout. - -Official source: - https://github.com/SwinTransformer/storage/releases/download/v1.0.2/cascade_mask_rcnn_swin_base_patch4_window7.pth - -Only the Swin backbone keys need remapping (layers->stages, attn/mlp naming). -neck / rpn_head / roi_head keys already match mmdet 3 CascadeRCNN. -""" - -from __future__ import annotations - -import argparse -import os -import sys - -import torch - -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) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument( - "--src", - default=os.path.join(REPO_ROOT, "checkpoints/task_networks/detection/model.pth"), - ) - parser.add_argument( - "--dst", - default=os.path.join(REPO_ROOT, "checkpoints/task_networks/detection/model_mmdet3.pth"), - ) - args = parser.parse_args() - - from mmdet.models.backbones.swin import swin_converter - - raw = torch.load(args.src, map_location="cpu") - state = raw["state_dict"] if isinstance(raw, dict) and "state_dict" in raw else raw - - backbone = {} - rest = {} - for k, v in state.items(): - if k.startswith("backbone."): - backbone[k[len("backbone.") :]] = v - else: - rest[k] = v - - converted_backbone = swin_converter(backbone) # adds 'backbone.' prefix - out_state = {} - out_state.update(converted_backbone) - out_state.update(rest) - - os.makedirs(os.path.dirname(args.dst), exist_ok=True) - torch.save({"state_dict": out_state, "meta": {"converted_from": args.src}}, args.dst) - print(f"Wrote {args.dst}") - print(f" backbone keys: {sum(1 for k in out_state if k.startswith('backbone.'))}") - print(f" other keys: {sum(1 for k in out_state if not k.startswith('backbone.'))}") - print(f" sample backbone: {[k for k in out_state if k.startswith('backbone.')][:5]}") - - -if __name__ == "__main__": - main() diff --git a/scripts/download_base_codecs.sh b/scripts/download_base_codecs.sh deleted file mode 100755 index 3b63b0a9052c376e73c7fcc79f08c6e7f77f8e06..0000000000000000000000000000000000000000 --- a/scripts/download_base_codecs.sh +++ /dev/null @@ -1,22 +0,0 @@ -#!/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 deleted file mode 100755 index 0af3b20766d3f52886af4a8e691faf4f4d1176a7..0000000000000000000000000000000000000000 --- a/scripts/eval_ctaic.py +++ /dev/null @@ -1,230 +0,0 @@ -#!/usr/bin/env python3 -"""Test / eval for C-TAIC: codec stats + optional full task-network metrics. - -Examples: - python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml - python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --with-metrics - python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --no-condition --with-metrics -""" - -from __future__ import annotations - -import argparse -import json -import os -import sys -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, ImageFolderDataset, build_test_transform -from flexicm.data.coco_eval import COCOEvalDataset, TASK_ANN_FILES, coco_eval_collate -from flexicm.models import CTAIC, TAIC -from flexicm.tasks import TASK_META, build_teacher -from flexicm.tasks.losses import TAICCriterion -from flexicm.tasks.metric_eval import run_task_metric_eval -from flexicm.tasks.metric_runners import ( - DEFAULT_TASK_NET_CKPTS, - DEFAULT_TASK_NET_CONFIGS, - build_metric_runner, -) -from flexicm.utils.codec_test import resolve_ckpt, test_ctaic_loader -from flexicm.utils.train_utils import load_checkpoint_dict, load_yaml_config, set_seed - -SCENARIOS = { - "s1": {"base": "detection", "ext": "instance"}, - "s2": {"base": "semantic", "ext": "panoptic"}, - "s3": {"base": "detection", "ext": "pose"}, -} - - -def parse_args(argv): - parser = argparse.ArgumentParser("Test FlexICM C-TAIC (codec + optional task metrics)") - parser.add_argument("-c", "--config", required=True) - given, remaining = parser.parse_known_args(argv) - cfg_path = given.config if os.path.isabs(given.config) else os.path.join(REPO_ROOT, given.config) - cfg = load_yaml_config(cfg_path) - # -c already consumed by the first parse; keep it as a default for the second pass - parser.set_defaults(config=cfg_path, **cfg) - for action in parser._actions: - if "--config" in action.option_strings: - action.required = False - break - parser.add_argument("--actual-bpp", action="store_true") - parser.add_argument("--no-condition", action="store_true") - parser.add_argument("--with-metrics", action="store_true") - parser.add_argument("--max-batches", type=int, default=None) - parser.add_argument("--split", type=str, default=None) - args = parser.parse_args(remaining) - args.config = cfg_path - if "--actual-bpp" in argv: - args.actual_bpp = True - if "--no-condition" in argv: - args.no_condition = True - if "--with-metrics" in argv: - args.with_metrics = True - return args - - -def build_codec_loader(args, ext_task, device): - split = args.split or getattr(args, "split", None) or "val2017" - tf = build_test_transform() - root = args.dataset_path - split_dir = os.path.join(root, split) - if os.path.isdir(split_dir): - if ext_task == "pose": - dataset = COCOWholeBodyImageDataset(root, split, tf) - else: - dataset = COCOImageDataset(root, split, tf) - else: - dataset = ImageFolderDataset(root, tf) - return DataLoader( - dataset, - batch_size=getattr(args, "test_batch_size", 1), - shuffle=False, - num_workers=getattr(args, "num_workers", 4), - pin_memory=(device == "cuda"), - ) - - -def build_metric_loader(args, ext_task, device): - split = args.split or getattr(args, "split", None) or "val2017" - ann_rel = getattr(args, "ann_file", None) or TASK_ANN_FILES.get(ext_task) - ann_file = ann_rel if os.path.isabs(ann_rel) else os.path.join(args.dataset_path, ann_rel) - dataset = COCOEvalDataset( - args.dataset_path, ann_file=ann_file, image_prefix=split, transform=build_test_transform() - ) - loader = DataLoader( - dataset, - batch_size=1, - shuffle=False, - num_workers=getattr(args, "num_workers", 4), - pin_memory=(device == "cuda"), - collate_fn=coco_eval_collate, - ) - return loader, ann_file - - -def main(argv): - args = parse_args(argv) - set_seed(getattr(args, "seed", 42)) - - os.environ["CUDA_VISIBLE_DEVICES"] = str(getattr(args, "gpu_id", 0)) - device = "cuda" if getattr(args, "cuda", True) and torch.cuda.is_available() else "cpu" - - scenario = args.scenario - base_task = SCENARIOS[scenario]["base"] - ext_task = SCENARIOS[scenario]["ext"] - base_meta = TASK_META[base_task] - ext_meta = TASK_META[ext_task] - out_channels = getattr(args, "out_channels", ext_meta["out_channels"]) - align_mode = getattr(args, "align_mode", ext_meta["align_mode"]) - lmbda = getattr(args, "lmbda", 0.0035) - use_condition = not bool(getattr(args, "no_condition", False)) - - ext_ckpt = resolve_ckpt(args.checkpoint, REPO_ROOT, label="C-TAIC checkpoint") - base_ckpt = resolve_ckpt(args.base_taic_checkpoint, REPO_ROOT, label="base TAIC checkpoint") - - print(f"Loading base TAIC ({base_task}): {base_ckpt}") - base = TAIC(N=128, M=192, out_channels=base_meta["out_channels"]).to(device) - state, _ = load_checkpoint_dict(base_ckpt, map_location=device) - base.load_state_dict(state, strict=False) - base.eval() - for p in base.parameters(): - p.requires_grad = False - - print(f"Loading C-TAIC extension ({ext_task}): {ext_ckpt}") - net = CTAIC(N=128, M=192, out_channels=out_channels).to(device) - state, _ = load_checkpoint_dict(ext_ckpt, map_location=device) - net.load_state_dict(state, strict=False) - net.eval() - - teacher = build_teacher( - ext_task, - pretrained_backbone=getattr(args, "pretrained_backbone", True), - use_official_teacher=getattr(args, "use_official_teacher", True), - task_config=getattr(args, "task_config", None), - task_checkpoint=getattr(args, "task_checkpoint", None), - device=device, - ) - teacher = teacher.to(device).eval() - criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode) - codec_loader = build_codec_loader(args, ext_task, device) - - codec_result = test_ctaic_loader( - net, - base, - teacher, - codec_loader, - criterion, - device, - use_condition=use_condition, - run_actual_bpp=bool(getattr(args, "actual_bpp", False)), - max_batches=args.max_batches, - ) - print("==== C-TAIC codec test summary ====") - for k, v in codec_result.items(): - print(f" {k}: {v:.6f}" if isinstance(v, float) else f" {k}: {v}") - - payload = { - "scenario": scenario, - "base_task": base_task, - "ext_task": ext_task, - "checkpoint": ext_ckpt, - "base_taic_checkpoint": base_ckpt, - "config": args.config, - "codec_result": codec_result, - } - - if getattr(args, "with_metrics", False): - task_cfg = getattr(args, "task_config", None) or DEFAULT_TASK_NET_CONFIGS[ext_task] - task_ckpt = getattr(args, "task_checkpoint", None) or DEFAULT_TASK_NET_CKPTS[ext_task] - if not os.path.isabs(task_cfg): - task_cfg = os.path.join(REPO_ROOT, task_cfg) - task_ckpt = resolve_ckpt(task_ckpt, REPO_ROOT, label=f"{ext_task} task-network checkpoint") - - print(f"[metric] loading extension task network:\n config={task_cfg}\n ckpt={task_ckpt}") - runner = build_metric_runner(ext_task, device=device) - runner.load(task_cfg, task_ckpt) - - metric_loader, ann_file = build_metric_loader(args, ext_task, device) - metrics = run_task_metric_eval( - net, - runner, - metric_loader, - device, - ann_file=ann_file, - use_condition=use_condition, - base_codec=base if use_condition else None, - max_batches=args.max_batches, - finalize_kwargs={ - "gt_folder": getattr(args, "panoptic_gt_folder", None), - "pred_folder": getattr(args, "panoptic_pred_folder", None), - "num_classes": getattr(args, "num_classes", 133), - }, - ) - print("==== C-TAIC task metric summary ====") - for k, v in metrics.items(): - print(f" {k}: {v}") - payload["task_config"] = task_cfg - payload["task_checkpoint"] = task_ckpt - payload["task_metrics"] = metrics - - out_dir = getattr(args, "result_dir", None) or os.path.join( - REPO_ROOT, "logs", "eval_ctaic", scenario, str(getattr(args, "quality_level", 1)) - ) - os.makedirs(out_dir, exist_ok=True) - out_json = os.path.join(out_dir, f"eval_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json") - with open(out_json, "w") as f: - json.dump(payload, f, indent=2, default=str) - print(f"Wrote {out_json}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/eval_taic.py b/scripts/eval_taic.py deleted file mode 100755 index 02d05b0b882154ece6b3ea00dbb01db9ecc94dd9..0000000000000000000000000000000000000000 --- a/scripts/eval_taic.py +++ /dev/null @@ -1,236 +0,0 @@ -#!/usr/bin/env python3 -"""Test / eval for TAIC: codec stats + optional full task-network metrics. - -Codec-only (default): - bpp, feature distortion D, loss; optional --actual-bpp - -With task metrics (--with-metrics): - also load official task-network config/checkpoint, run truncated task net from h, - report mAP-bbox / mAP-mask / mIoU / PQ / mAP-OKS (task-dependent). - -Examples: - python scripts/eval_taic.py -c configs/eval/taic_detection.yaml - python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --with-metrics - python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --with-metrics --max-batches 50 -""" - -from __future__ import annotations - -import argparse -import json -import os -import sys -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, - ImageFolderDataset, - build_test_transform, -) -from flexicm.data.coco_eval import COCOEvalDataset, TASK_ANN_FILES, coco_eval_collate -from flexicm.models import TAIC -from flexicm.tasks import TASK_META, build_teacher -from flexicm.tasks.losses import TAICCriterion -from flexicm.tasks.metric_eval import run_task_metric_eval -from flexicm.tasks.metric_runners import ( - DEFAULT_TASK_NET_CKPTS, - DEFAULT_TASK_NET_CONFIGS, - build_metric_runner, -) -from flexicm.utils.codec_test import resolve_ckpt, test_taic_loader -from flexicm.utils.train_utils import load_checkpoint_dict, load_yaml_config, set_seed - - -def parse_args(argv): - parser = argparse.ArgumentParser("Test FlexICM TAIC (codec + optional task metrics)") - parser.add_argument("-c", "--config", required=True, help="configs/eval/taic_*.yaml") - given, remaining = parser.parse_known_args(argv) - cfg_path = given.config if os.path.isabs(given.config) else os.path.join(REPO_ROOT, given.config) - cfg = load_yaml_config(cfg_path) - # -c already consumed by the first parse; keep it as a default for the second pass - parser.set_defaults(config=cfg_path, **cfg) - for action in parser._actions: - if "--config" in action.option_strings: - action.required = False - break - parser.add_argument("--actual-bpp", action="store_true") - parser.add_argument("--with-metrics", action="store_true", help="Run full task-network metrics") - parser.add_argument("--max-batches", type=int, default=None) - parser.add_argument("--split", type=str, default=None) - parser.add_argument( - "--eval-size", - type=int, - default=cfg.get("eval_size", None), - help="If set (e.g. 256), resize eval images to eval_size×eval_size", - ) - args = parser.parse_args(remaining) - args.config = cfg_path - if "--actual-bpp" in argv: - args.actual_bpp = True - if "--with-metrics" in argv: - args.with_metrics = True - return args - - -def build_codec_loader(args, device): - split = args.split or getattr(args, "split", None) or "val2017" - eval_size = getattr(args, "eval_size", None) - tf = build_test_transform(eval_size=eval_size) - root = args.dataset_path - split_dir = os.path.join(root, split) - if os.path.isdir(split_dir): - if args.task == "pose": - dataset = COCOWholeBodyImageDataset(root, split, tf) - else: - dataset = COCOImageDataset(root, split, tf) - else: - dataset = ImageFolderDataset(root, tf) - return DataLoader( - dataset, - batch_size=getattr(args, "test_batch_size", 1), - shuffle=False, - num_workers=getattr(args, "num_workers", 4), - pin_memory=(device == "cuda"), - ) - - -def build_metric_loader(args, device): - split = args.split or getattr(args, "split", None) or "val2017" - ann_rel = getattr(args, "ann_file", None) or TASK_ANN_FILES.get(args.task) - if ann_rel is None: - raise ValueError(f"No ann_file for task={args.task}") - ann_file = ann_rel if os.path.isabs(ann_rel) else os.path.join(args.dataset_path, ann_rel) - eval_size = getattr(args, "eval_size", None) - dataset = COCOEvalDataset( - args.dataset_path, - ann_file=ann_file, - image_prefix=split, - transform=build_test_transform(eval_size=eval_size), - ) - loader = DataLoader( - dataset, - batch_size=1, - shuffle=False, - num_workers=getattr(args, "num_workers", 4), - pin_memory=(device == "cuda"), - collate_fn=coco_eval_collate, - ) - return loader, ann_file - - -def main(argv): - args = parse_args(argv) - set_seed(getattr(args, "seed", 42)) - - os.environ["CUDA_VISIBLE_DEVICES"] = str(getattr(args, "gpu_id", 0)) - device = "cuda" if getattr(args, "cuda", True) 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"]) - lmbda = getattr(args, "lmbda", 0.0035) - - ckpt = resolve_ckpt(args.checkpoint, REPO_ROOT, label="TAIC checkpoint") - print(f"Loading TAIC checkpoint: {ckpt}") - - net = TAIC(N=128, M=192, out_channels=out_channels).to(device) - state, _ = load_checkpoint_dict(ckpt, map_location=device) - missing = net.load_state_dict(state, strict=False) - print(f"load_state_dict: missing={len(missing.missing_keys)} unexpected={len(missing.unexpected_keys)}") - net.eval() - - # ---- codec test ---- - teacher = build_teacher( - task, - pretrained_backbone=getattr(args, "pretrained_backbone", True), - use_official_teacher=getattr(args, "use_official_teacher", True), - task_config=getattr(args, "task_config", None), - task_checkpoint=getattr(args, "task_checkpoint", None), - device=device, - ) - teacher = teacher.to(device).eval() - criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode) - codec_loader = build_codec_loader(args, device) - eval_size = getattr(args, "eval_size", None) - size_msg = f"{eval_size}x{eval_size}" if eval_size else "original" - print(f"[codec] test set size: {len(codec_loader.dataset)} device={device} task={task} input={size_msg}") - - codec_result = test_taic_loader( - net, - teacher, - codec_loader, - criterion, - device, - align_divisor=256, - run_actual_bpp=bool(getattr(args, "actual_bpp", False)), - max_batches=args.max_batches, - ) - print("==== TAIC codec test summary ====") - for k, v in codec_result.items(): - print(f" {k}: {v:.6f}" if isinstance(v, float) else f" {k}: {v}") - - payload = { - "task": task, - "checkpoint": ckpt, - "config": args.config, - "codec_result": codec_result, - } - - # ---- optional task metrics ---- - if getattr(args, "with_metrics", False): - task_cfg = getattr(args, "task_config", None) or DEFAULT_TASK_NET_CONFIGS[task] - task_ckpt = getattr(args, "task_checkpoint", None) or DEFAULT_TASK_NET_CKPTS[task] - if not os.path.isabs(task_cfg): - task_cfg = os.path.join(REPO_ROOT, task_cfg) - task_ckpt = resolve_ckpt(task_ckpt, REPO_ROOT, label=f"{task} task-network checkpoint") - - print(f"[metric] loading task network:\n config={task_cfg}\n ckpt={task_ckpt}") - runner = build_metric_runner(task, device=device) - runner.load(task_cfg, task_ckpt) - - metric_loader, ann_file = build_metric_loader(args, device) - print(f"[metric] COCO eval images: {len(metric_loader.dataset)} ann={ann_file}") - metrics = run_task_metric_eval( - net, - runner, - metric_loader, - device, - ann_file=ann_file, - use_condition=False, - base_codec=None, - max_batches=args.max_batches, - finalize_kwargs={ - "gt_folder": getattr(args, "panoptic_gt_folder", None), - "pred_folder": getattr(args, "panoptic_pred_folder", None), - "num_classes": getattr(args, "num_classes", 133), - }, - ) - print("==== TAIC task metric summary ====") - for k, v in metrics.items(): - print(f" {k}: {v}") - payload["task_config"] = task_cfg - payload["task_checkpoint"] = task_ckpt - payload["task_metrics"] = metrics - - out_dir = getattr(args, "result_dir", None) or os.path.join( - REPO_ROOT, "logs", "eval_taic", task, str(getattr(args, "quality_level", 1)) - ) - os.makedirs(out_dir, exist_ok=True) - out_json = os.path.join(out_dir, f"eval_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json") - with open(out_json, "w") as f: - json.dump(payload, f, indent=2, default=str) - print(f"Wrote {out_json}") - return 0 - - -if __name__ == "__main__": - raise SystemExit(main(sys.argv[1:])) diff --git a/scripts/sanity_check.py b/scripts/sanity_check.py deleted file mode 100644 index d9d1bdfa2651575841ac8d5665bcb07b6a96fa24..0000000000000000000000000000000000000000 --- a/scripts/sanity_check.py +++ /dev/null @@ -1,31 +0,0 @@ -#!/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 deleted file mode 100755 index e3b9701767f44c4c97ddb47a727d5a6a81756c3c..0000000000000000000000000000000000000000 --- a/scripts/train_ctaic.py +++ /dev/null @@ -1,253 +0,0 @@ -#!/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.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, -) - -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) - # -c/--stage already consumed by the first parse; keep them for the second pass - parser.set_defaults(config=given.config, stage=given.stage, **cfg) - for action in parser._actions: - if "--config" in action.option_strings: - action.required = False - break - 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, align_divisor=256): - ext_model.eval() - meters = {k: AverageMeter() for k in ("loss", "bpp", "distortion")} - for images in loader: - images = images.to(device) - # Pad to codec/Swin-friendly size (same as train_taic.validate). - align = Alignment(divisor=align_divisor, mode="pad", padding_mode="constant").to(device) - x = align.align(images) - if stage == 1: - out = ext_model(x, use_condition=False) - else: - y_b = encode_base_latent(base_model, x) - out = ext_model(x, y_b_hat=y_b, use_condition=True) - gt = teacher.gt_features(x) - 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), - use_official_teacher=getattr(args, "use_official_teacher", True), - task_config=getattr(args, "task_config", None), - task_checkpoint=getattr(args, "task_checkpoint", None), - device=device, - ) - 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 deleted file mode 100755 index 45a7a2632eb19cf1d9a9781c0699b4d1421839ae..0000000000000000000000000000000000000000 --- a/scripts/train_taic.py +++ /dev/null @@ -1,233 +0,0 @@ -#!/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) - # -c already consumed by the first parse; keep it as a default for the second pass - parser.set_defaults(config=given.config, **cfg) - for action in parser._actions: - if "--config" in action.option_strings: - action.required = False - break - 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) - # Keep teacher + codec features on the aligned (padded) grid so the Swin - # backbone always sees a patch/window-divisible size. bpp still uses the - # original pixel count below. - with torch.no_grad(): - gt = teacher.gt_features(x) - 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()) - 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, "val2017", train_tf) - else: - train_set = COCOImageDataset(args.dataset_path, "val2017", 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) - if getattr(args, "freeze_base_codec", True): - net.freeze_base_codec() - logging.info("Base TIC codec frozen") - else: - for p in net.parameters(): - p.requires_grad = True - logging.info("Base TIC codec unfrozen; training the entire TAIC model") - 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), - use_official_teacher=getattr(args, "use_official_teacher", True), - task_config=getattr(args, "task_config", None), - task_checkpoint=getattr(args, "task_checkpoint", None), - device=device, - ) - 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) - if epoch % 2 == 0: - 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:])