Add task-metric eval path and Cascade Mask R-CNN + Swin-B for det/instance.

Wire --with-metrics for TAIC/C-TAIC, point detection and instance to the official Cascade Mask R-CNN Swin-B zoo (out_channels=128), and document task-network setup for server runs.

Co-authored-by: Cursor <cursoragent@cursor.com>

README.md

@@ -5,36 +5,46 @@ Official codebase for the paper **FlexICM: A Flexible Image Coding for Machines
 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
+- **C-TAIC (Extension Layer)**: three multi-task scenarios that condition on the base-layer latent \hat{y}_b via cross-attention
 
 ## Five Tasks and Three Scenarios
 
 ### TAIC (five task codecs)
 
-| Task | Teacher / Task Network | Feature Alignment | Metric |
-|------|------------------------|-------------------|--------|
-| Object Detection | Faster R-CNN + **Swin-B** | FPN `P2..P6` (Eq. 2) | mAP-bbox |
-| Semantic Segmentation | UPerNet + **Swin-B** | FPN `P2..P6` | mIoU |
-| Instance Segmentation | Mask R-CNN + **Swin-B** | FPN `P2..P6` | mAP-mask |
-| Panoptic Segmentation | MaskFormer + **Swin-B** | Stages `F1..F4` (Eq. 3) | PQ |
-| Pose Estimation | **HigherHRNet** | Stages `F1..F4` | mAP-OKS |
+
+| 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 |
+
+| 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
+> 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)
@@ -48,13 +58,19 @@ 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 |
+
+| 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**
 
@@ -62,22 +78,25 @@ Teachers are already implemented in `flexicm/tasks/` and are constructed automat
 `scripts/train_taic.py` / `scripts/train_ctaic.py` via `build_teacher(...)`.
 You still must install their runtime dependencies and allow pretrained weights to download.
 
-| Task | Teacher used in training | What you need installed |
-|------|--------------------------|-------------------------|
-| Detection / Instance / Semantic / Panoptic | Swin-B backbone (+ FPN or stages) via `timm` | `timm` (from `requirements.txt`); first run downloads ImageNet-pretrained Swin-B |
-| Pose | HigherHRNet-style HRNet stem (original HRNet, not Swin) | Implemented in-repo; no extra package beyond PyTorch |
+
+| 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-B weights (for the four Swin tasks)
+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\).
+Without a working teacher, training will fail when computing the feature-alignment term D.
 
 ### Task heads for metric evaluation
 
@@ -87,12 +106,13 @@ To evaluate paper metrics (mAP / mIoU / PQ / OKS) with full task heads, also ins
 pip install -U openmim
 mim install mmengine mmcv
 mim install mmdet mmsegmentation mmpose
-# or Detectron2 (alternative for detection / instance evaluation)
 ```
 
-Recommended official weights (same model families as the paper):
+Official detection / instance weights from
+[Swin-Transformer-Object-Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection)
+(see `configs/task_networks/README.md`):
 
-- **Faster / Mask R-CNN + Swin-B**: MMDetection Model Zoo
+- **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**)
@@ -101,6 +121,8 @@ These full heads are **not** required to start codec training; they are for fina
 
 ---
 
+
+
 ## Repository Layout
 
 ```
@@ -139,8 +161,12 @@ Eval / codec-test configs: `configs/eval/`.
 
 ---
 
+
+
 ## Dataset Preparation
 
+
+
 ### COCO-2017 (detection / instance / semantic / panoptic)
 
 ```text
@@ -175,6 +201,8 @@ Set in the corresponding YAML:
 dataset_path: "/data/coco2017"
 ```
 
+
+
 ### COCO-WholeBody (pose estimation)
 
 Pose uses the same COCO `train2017/val2017` images plus WholeBody keypoint annotations:
@@ -201,16 +229,20 @@ Edit points: `flexicm/data/datasets.py`, `flexicm/tasks/swin_teacher.py`, `flexi
 
 ---
 
+
+
 ## 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) |
+
+| 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
@@ -231,9 +263,10 @@ Trained TAIC / C-TAIC weights for eval should be placed under `checkpoints/taic/
 `checkpoints/ctaic/` (see `checkpoints/README.md`). Until then, each quality folder
 contains a `PLACEHOLDER` file.
 
-
 ---
 
+
+
 ## Training
 
 Paper settings:
@@ -245,6 +278,8 @@ Paper settings:
 
 > If GPU memory is insufficient, reduce `batch_size` (optionally use gradient accumulation to approximate the paper effective batch).
 
+
+
 ### Train five TAIC models
 
 ```bash
@@ -260,7 +295,7 @@ 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.
+Requires a trained **base TAIC** checkpoint (to provide \hat{y}_b) and Stage-1 weights for the extension task.
 
 ```bash
 # ---- s1: det → instance ----
@@ -278,10 +313,12 @@ 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 |
+
+| 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:
 
@@ -293,64 +330,113 @@ stage1_checkpoint:     # Stage-1 result loaded in Stage 2
 
 ---
 
+
+
 ## Codec Test
 
+Codec test measures **compression statistics**:
 
-For C-TAIC, reported `bpp` is **extension-layer only** (base-layer rate is excluded), matching the paper.
 
-### Prepare checkpoints
+| 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()` |
 
-1. Train models (or copy trained weights) into the `checkpoints/` tree — Download checkpoints.
-2. Edit `dataset_path` / `gpu_id` in `configs/eval/*.yaml`
 
-### Test TAIC (five tasks)
+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_taic.py -c configs/eval/taic_semantic.yaml
-python scripts/eval_taic.py -c configs/eval/taic_instance.yaml
-python scripts/eval_taic.py -c configs/eval/taic_panoptic.yaml
-python scripts/eval_taic.py -c configs/eval/taic_pose.yaml
-
-# optional: also measure actual entropy-coded bitstream bpp
 python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --actual-bpp
-
-# optional: smoke test on a few batches
 python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --max-batches 10
+
+python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml
+python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --no-condition
 ```
 
-Results JSON is written under `logs/eval_taic/<task>/<quality>/`.
+---
+
+
+
+## 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:
+
 
-### Test C-TAIC (three scenarios)
+| 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
-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
+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
+```
 
-# disable base-layer conditioning (TAIC-mode / graceful degradation)
-python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --no-condition
 
-python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --actual-bpp
+
+### 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
 ```
 
-Results JSON is written under `logs/eval_ctaic/<scenario>/<quality>/`.
+1. Set in each `configs/eval/*.yaml`:
 
-> Task rate–accuracy curves (mAP / mIoU / PQ / OKS vs bpp) are **not** included yet.
+```yaml
+task_config: "./configs/task_networks/<real_config>.py"
+task_checkpoint: "./checkpoints/task_networks/<task>/model.pth"
+ann_file: "annotations/instances_val2017.json"
+```
 
----
 
-## Code Map to the Paper
 
-| Paper component | Code location |
-|-----------------|---------------|
-| SFMA | `flexicm/models/sfma.py` |
-| Task Connector | `flexicm/models/task_connector.py` |
-| TAIC | `flexicm/models/taic.py` |
-| C-TAIC + two-stage freeze | `flexicm/models/ctaic.py` |
-| Prompt / Mask / Cd | `flexicm/models/conditional.py` |
-| Cross-attention (Q from features; K/V include prompts) | `flexicm/models/cross_attention.py` |
-| \(R+\lambda D\) | `flexicm/tasks/losses.py` |
-| Five teachers | `flexicm/tasks/__init__.py` |
+### Run codec + metrics
+
+```bash
+python scripts/eval_taic.py -c configs/eval/taic_detection.yaml --with-metrics
+python scripts/eval_taic.py -c configs/eval/taic_instance.yaml --with-metrics
+python scripts/eval_taic.py -c configs/eval/taic_semantic.yaml --with-metrics
+python scripts/eval_taic.py -c configs/eval/taic_panoptic.yaml --with-metrics
+python scripts/eval_taic.py -c configs/eval/taic_pose.yaml --with-metrics
+
+python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --with-metrics
+python scripts/eval_ctaic.py -c configs/eval/ctaic_s2.yaml --with-metrics
+python scripts/eval_ctaic.py -c configs/eval/ctaic_s3.yaml --with-metrics
+```
+
+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.
 

checkpoints/README.md

@@ -70,3 +70,23 @@ 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.

checkpoints/task_networks/detection/PLACEHOLDER

@@ -0,0 +1,2 @@
+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.

checkpoints/task_networks/instance/PLACEHOLDER

@@ -0,0 +1,2 @@
+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.

checkpoints/task_networks/panoptic/PLACEHOLDER

@@ -0,0 +1,2 @@
+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.

checkpoints/task_networks/pose/PLACEHOLDER

@@ -0,0 +1,2 @@
+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.

checkpoints/task_networks/semantic/PLACEHOLDER

@@ -0,0 +1,2 @@
+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.

configs/ctaic/s1_det_instance.yaml

@@ -22,3 +22,5 @@ cuda: true
 save: true
 seed: 42
 pretrained_backbone: true
+out_channels: 128  # extension = Cascade Mask R-CNN + Swin-B
+align_mode: "fpn"

configs/eval/ctaic_s1.yaml

@@ -1,12 +1,14 @@
-# Codec test config for C-TAIC scenario s1 (det -> instance)
 scenario: "s1"
 dataset_path: "/data/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"
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/instance/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/ctaic_s2.yaml

@@ -1,11 +1,15 @@
 scenario: "s2"
 dataset_path: "/data/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"
+task_config: "./configs/task_networks/maskformer_swin-b_coco.py"
+task_checkpoint: "./checkpoints/task_networks/panoptic/model.pth"
+panoptic_gt_folder: "/data/coco2017/annotations/panoptic_val2017"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/ctaic_s3.yaml

@@ -1,11 +1,14 @@
 scenario: "s3"
 dataset_path: "/data/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"
+task_config: "./configs/task_networks/higherhrnet_w32_coco_wholebody.py"
+task_checkpoint: "./checkpoints/task_networks/pose/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/taic_detection.yaml

@@ -1,12 +1,15 @@
-# Codec test config for TAIC / detection
-# Place real weights at checkpoint path (replace PLACEHOLDER).
+# TAIC eval: codec test + optional --with-metrics
 task: "detection"
 dataset_path: "/data/coco2017"
 split: "val2017"
+ann_file: "annotations/instances_val2017.json"
 quality_level: 1
 lmbda: 0.0035
 checkpoint: "./checkpoints/taic/detection/1/checkpoint_best_loss.pth.tar"
 base_codec: "./checkpoints/base_codec/base_codec_1.pth.tar"
+# official task network (required for --with-metrics)
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/detection/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/taic_instance.yaml

@@ -1,12 +1,15 @@
-# Codec test config for TAIC / instance
-# Place real weights at checkpoint path (replace PLACEHOLDER).
+# TAIC eval: codec test + optional --with-metrics
 task: "instance"
 dataset_path: "/data/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)
+task_config: "./configs/task_networks/cascade_mask_rcnn_swin_base_coco.py"
+task_checkpoint: "./checkpoints/task_networks/instance/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/taic_panoptic.yaml

@@ -1,12 +1,15 @@
-# Codec test config for TAIC / panoptic
-# Place real weights at checkpoint path (replace PLACEHOLDER).
+# TAIC eval: codec test + optional --with-metrics
 task: "panoptic"
 dataset_path: "/data/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)
+task_config: "./configs/task_networks/maskformer_swin-b_coco.py"
+task_checkpoint: "./checkpoints/task_networks/panoptic/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/taic_pose.yaml

@@ -1,12 +1,15 @@
-# Codec test config for TAIC / pose
-# Place real weights at checkpoint path (replace PLACEHOLDER).
+# TAIC eval: codec test + optional --with-metrics
 task: "pose"
 dataset_path: "/data/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)
+task_config: "./configs/task_networks/higherhrnet_w32_coco_wholebody.py"
+task_checkpoint: "./checkpoints/task_networks/pose/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/eval/taic_semantic.yaml

@@ -1,12 +1,15 @@
-# Codec test config for TAIC / semantic
-# Place real weights at checkpoint path (replace PLACEHOLDER).
+# TAIC eval: codec test + optional --with-metrics
 task: "semantic"
 dataset_path: "/data/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)
+task_config: "./configs/task_networks/upernet_swin-b_coco.py"
+task_checkpoint: "./checkpoints/task_networks/semantic/model.pth"
 gpu_id: 0
 cuda: true
 test_batch_size: 1

configs/taic/detection.yaml

@@ -1,4 +1,4 @@
-# FlexICM TAIC - Object Detection (Faster R-CNN + Swin-B)
+# FlexICM TAIC - Object Detection (Cascade Mask R-CNN + Swin-B)
 root: "logs"
 exp_name: "taic_detection"
 task: "detection"

configs/taic/instance.yaml

@@ -1,3 +1,4 @@
+# FlexICM TAIC - Instance Segmentation (Cascade Mask R-CNN + Swin-B)
 root: "logs"
 exp_name: "taic_instance"
 task: "instance"
@@ -17,5 +18,5 @@ cuda: true
 save: true
 seed: 42
 pretrained_backbone: true
-out_channels: 128
+out_channels: 128  # Cascade Mask R-CNN + Swin-B F1
 align_mode: "fpn"

configs/task_networks/README.md

@@ -0,0 +1,44 @@
+# 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`.

configs/task_networks/cascade_mask_rcnn_swin_base_coco.py

@@ -0,0 +1,19 @@
+# STUB / pointer config for Cascade Mask R-CNN + Swin-B (official Swin detection zoo).
+#
+# Source repo: https://github.com/SwinTransformer/Swin-Transformer-Object-Detection
+# Official config:
+#   configs/swin/cascade_mask_rcnn_swin_base_patch4_window7_mstrain_480-800_giou_4conv1f_adamw_3x_coco.py
+# Official checkpoint:
+#   https://github.com/SwinTransformer/storage/releases/download/v1.0.2/cascade_mask_rcnn_swin_base_patch4_window7.pth
+#
+# Replace this file by copying/symlinking the real config from that repo, then set
+# task_checkpoint to the downloaded .pth under checkpoints/task_networks/detection/.
+#
+# Used for FlexICM detection (mAP-bbox) and instance segmentation (mAP-mask).
+# Backbone F1 channels = 128 (Swin-B).
+
+raise RuntimeError(
+    "Replace configs/task_networks/cascade_mask_rcnn_swin_base_coco.py with the official "
+    "Swin-Transformer-Object-Detection config: "
+    "cascade_mask_rcnn_swin_base_patch4_window7_mstrain_480-800_giou_4conv1f_adamw_3x_coco.py"
+)

configs/task_networks/higherhrnet_w32_coco_wholebody.py

@@ -0,0 +1,7 @@
+# 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)."
+)

configs/task_networks/maskformer_swin-b_coco.py

@@ -0,0 +1,7 @@
+# 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)."
+)

configs/task_networks/upernet_swin-b_coco.py

@@ -0,0 +1,7 @@
+# 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)."
+)

flexicm/data/__init__.py

@@ -6,12 +6,16 @@ from .datasets import (
     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",
 ]

flexicm/data/coco_eval.py

@@ -0,0 +1,80 @@
+"""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",
+}

flexicm/tasks/__init__.py

@@ -1,8 +1,8 @@
 """Task-specific frozen teachers for FlexICM feature alignment.
 
 Five tasks (paper Sec.III.A / IV.A):
-  1. Object detection      - Faster R-CNN + Swin-B (FPN P2-P6)
-  2. Instance segmentation - Mask R-CNN + Swin-B (FPN P2-P6)
+  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)
@@ -24,15 +24,22 @@ from flexicm.tasks.swin_teacher import SwinStageTeacher
 
 
 class DetectionTeacher(nn.Module):
-    """Faster R-CNN / Mask R-CNN style: align FPN P2..P6."""
+    """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  # Swin-B F1
+    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)
+        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]:
@@ -194,7 +201,7 @@ TASK_META = {
     },
     "instance": {
         "align_mode": "fpn",
-        "out_channels": 128,
+        "out_channels": 128,  # Cascade Mask R-CNN + Swin-B F1
         "metric": "mAP-mask",
         "dataset": "coco",
     },

flexicm/tasks/metric_eval.py

@@ -0,0 +1,84 @@
+"""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
+            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 = crop_feature_to_image(out["h"], (H, W))
+            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

flexicm/tasks/metric_runners.py

@@ -0,0 +1,479 @@
+"""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/MMSeg SwinTransformer that exposes `.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")
+
+    outs = [h]
+    x = h
+    # Stage 0 already produced F1 (=h); run stages 1..N-1
+    for i in range(1, len(stages)):
+        x = stages[i](x)
+        if isinstance(x, (tuple, list)):
+            x = x[0]
+        # MMDet Swin may return NCHW already
+        if x.dim() == 4 and x.shape[1] < x.shape[-1] and x.shape[-1] in (96, 128, 192, 256, 384, 512, 768, 1024):
+            # likely NHWC
+            x = x.permute(0, 3, 1, 2).contiguous()
+        outs.append(x)
+
+    # Apply per-stage norms if present (mmdet Swin)
+    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
+                # LayerNorm over channel last in some impls
+                norm = getattr(backbone, norm_name)
+                if nchw.shape[1] == getattr(norm, "normalized_shape", [nchw.shape[1]])[0] if hasattr(norm, "normalized_shape") else True:
+                    # try NCHW LayerNorm via transpose
+                    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)
+        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
+        H, W = int(img_meta["height"]), int(img_meta["width"])
+        try:
+            # MMDet 3.x style
+            from mmdet.structures import DetDataSample
+            from mmengine.structures import InstanceData
+
+            data_sample = DetDataSample()
+            data_sample.set_metainfo(
+                dict(
+                    img_shape=(H, W),
+                    ori_shape=(H, W),
+                    pad_shape=(H, W),
+                    scale_factor=(1.0, 1.0),
+                    img_id=img_meta.get("image_id"),
+                )
+            )
+            # Use RPN + ROI heads with injected feats
+            if hasattr(self.model, "extract_feat"):
+                # bypass extract_feat by calling predict with feats if supported
+                pass
+            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]
+            inst = pred.pred_instances
+            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=(H, W, 3),
+                    ori_shape=(H, W, 3),
+                    pad_shape=(H, W, 3),
+                    scale_factor=1.0,
+                    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",
+}

flexicm/tasks/swin_teacher.py

@@ -1,6 +1,8 @@
-"""Swin-B backbone helpers shared by detection / segmentation teachers.
+"""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.
 
-Matches Fig.1(c): Stage depths [2,2,18,2], F1 at H/4 with C=128 (Swin-B).
 h from TAIC replaces F1 and is fed into Stage 2 onward.
 """
 
@@ -14,6 +16,12 @@ 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)."""
@@ -36,45 +44,58 @@ class SimpleFPN(nn.Module):
         return {"p2": p2, "p3": p3, "p4": p4, "p5": p5, "p6": p6}
 
 
-def build_swin_b_backbone(pretrained: bool = True):
-    """Build Swin-B via timm; returns backbone module with forward_features stages."""
+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-B teachers: pip install timm") from 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)}")
 
-    # features_only gives list of stage outputs
     model = timm.create_model(
-        "swin_base_patch4_window7_224",
+        _SWIN_TIMM_NAMES[key],
         pretrained=pretrained,
         features_only=True,
         out_indices=(0, 1, 2, 3),
-        img_size=224,  # overridden dynamically by dynamic image size support in newer timm
+        img_size=224,
     )
     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-B backbone.
+    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):
+    def __init__(
+        self,
+        pretrained: bool = True,
+        use_fpn: bool = True,
+        fpn_dim: int = 256,
+        swin_variant: str = "base",
+    ):
         super().__init__()
-        self.backbone = freeze_module(build_swin_b_backbone(pretrained=pretrained))
-        # timm swin_base features_only channel dims
-        self.feat_channels = list(self.backbone.feature_info.channels())  # typically [128,256,512,1024]
+        self.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
 
-        # Build stage modules for truncated forward from F1.
-        # timm Swin features_only structure varies; we use a practical approach:
-        # full forward for GT; for truncated, interpolate/project h and run full backbone
-        # with early feature replacement via forward hooks when possible.
         self._f1_dim = self.feat_channels[0]
 
     @property

requirements.txt

@@ -7,12 +7,13 @@ Pillow>=8.0
 tqdm>=4.60
 numpy>=1.20
 einops>=0.4.0
+pycocotools>=2.0.6
 
-# Optional — full task-network evaluation (mAP / mIoU / PQ / OKS)
+# Optional — required for --with-metrics (full task-network evaluation)
 # openmim
 # mmengine
 # mmcv
 # mmdet
 # mmsegmentation
 # mmpose
-# detectron2
+# panopticapi   # for PQ

scripts/eval_ctaic.py

@@ -1,14 +1,10 @@
 #!/usr/bin/env python3
-"""Codec test for C-TAIC (extension layer).
+"""Test / eval for C-TAIC: codec stats + optional full task-network metrics.
 
-Reports extension-layer likelihood bpp and feature distortion D (paper: bpp excludes base layer).
-Optional actual bitstream bpp via compress/decompress.
-Does NOT compute task metrics (mAP / mIoU / PQ / OKS) — those come later.
-
-Example:
+Examples:
   python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml
-  python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --no-condition   # TAIC-mode ablation
-  python scripts/eval_ctaic.py -c configs/eval/ctaic_s1.yaml --actual-bpp
+  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
@@ -27,9 +23,16 @@ 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
 
@@ -41,14 +44,15 @@ SCENARIOS = {
 
 
 def parse_args(argv):
-    parser = argparse.ArgumentParser("Codec test: FlexICM C-TAIC")
-    parser.add_argument("-c", "--config", required=True, help="configs/eval/ctaic_*.yaml")
+    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)
     parser.set_defaults(**cfg)
     parser.add_argument("--actual-bpp", action="store_true")
-    parser.add_argument("--no-condition", action="store_true", help="Disable base-layer conditioning (TAIC mode)")
+    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)
@@ -57,10 +61,12 @@ def parse_args(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_loader(args, ext_task, device):
+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
@@ -81,6 +87,24 @@ def build_loader(args, ext_task, device):
     )
 
 
+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))
@@ -118,38 +142,23 @@ def main(argv):
     teacher = build_teacher(ext_task, pretrained_backbone=getattr(args, "pretrained_backbone", True))
     teacher = teacher.to(device).eval()
     criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode)
+    codec_loader = build_codec_loader(args, ext_task, device)
 
-    loader = build_loader(args, ext_task, device)
-    print(
-        f"Test set size: {len(loader.dataset)}  device={device}  "
-        f"scenario={scenario}  use_condition={use_condition}"
-    )
-
-    result = test_ctaic_loader(
+    codec_result = test_ctaic_loader(
         net,
         base,
         teacher,
-        loader,
+        codec_loader,
         criterion,
         device,
         use_condition=use_condition,
-        align_divisor=256,
         run_actual_bpp=bool(getattr(args, "actual_bpp", False)),
         max_batches=args.max_batches,
     )
-
     print("==== C-TAIC codec test summary ====")
-    for k, v in result.items():
-        if isinstance(v, float):
-            print(f"  {k}: {v:.6f}")
-        else:
-            print(f"  {k}: {v}")
+    for k, v in codec_result.items():
+        print(f"  {k}: {v:.6f}" if isinstance(v, float) else f"  {k}: {v}")
 
-    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"codec_test_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
     payload = {
         "scenario": scenario,
         "base_task": base_task,
@@ -157,10 +166,50 @@ def main(argv):
         "checkpoint": ext_ckpt,
         "base_taic_checkpoint": base_ckpt,
         "config": args.config,
-        "result": result,
+        "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)
+        json.dump(payload, f, indent=2, default=str)
     print(f"Wrote {out_json}")
     return 0
 

scripts/eval_taic.py

@@ -1,12 +1,17 @@
 #!/usr/bin/env python3
-"""Codec test for TAIC (base layer).
+"""Test / eval for TAIC: codec stats + optional full task-network metrics.
 
-Reports likelihood bpp, feature-alignment distortion D, and optional actual bitstream bpp.
-Does NOT compute task metrics (mAP / mIoU / PQ / OKS) — those come later.
+Codec-only (default):
+  bpp, feature distortion D, loss; optional --actual-bpp
 
-Example:
+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 --actual-bpp
+  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
@@ -30,31 +35,41 @@ from flexicm.data import (
     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("Codec test: FlexICM TAIC")
+    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)
     parser.set_defaults(**cfg)
-    parser.add_argument("--actual-bpp", action="store_true", help="Also run compress/decompress bpp")
-    parser.add_argument("--max-batches", type=int, default=None, help="Limit batches for a smoke test")
-    parser.add_argument("--split", type=str, default=None, help="Override image split folder, e.g. val2017")
+    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)
     args = parser.parse_args(remaining)
     args.config = cfg_path
-    if given.__dict__.get("actual_bpp") or "--actual-bpp" in argv:
+    if "--actual-bpp" in argv:
         args.actual_bpp = True
+    if "--with-metrics" in argv:
+        args.with_metrics = True
     return args
 
 
-def build_loader(args, device):
+def build_codec_loader(args, device):
     split = args.split or getattr(args, "split", None) or "val2017"
     tf = build_test_transform()
     root = args.dataset_path
@@ -66,7 +81,6 @@ def build_loader(args, device):
             dataset = COCOImageDataset(root, split, tf)
     else:
         dataset = ImageFolderDataset(root, tf)
-
     return DataLoader(
         dataset,
         batch_size=getattr(args, "test_batch_size", 1),
@@ -76,6 +90,29 @@ def build_loader(args, device):
     )
 
 
+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)
+    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))
@@ -98,44 +135,77 @@ def main(argv):
     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))
     teacher = teacher.to(device).eval()
     criterion = TAICCriterion(lmbda=lmbda, align_mode=align_mode)
+    codec_loader = build_codec_loader(args, device)
+    print(f"[codec] test set size: {len(codec_loader.dataset)}  device={device}  task={task}")
 
-    loader = build_loader(args, device)
-    print(f"Test set size: {len(loader.dataset)}  device={device}  task={task}")
-
-    result = test_taic_loader(
+    codec_result = test_taic_loader(
         net,
         teacher,
-        loader,
+        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 result.items():
-        if isinstance(v, float):
-            print(f"  {k}: {v:.6f}")
-        else:
-            print(f"  {k}: {v}")
+    for k, v in codec_result.items():
+        print(f"  {k}: {v:.6f}" if isinstance(v, float) else f"  {k}: {v}")
 
-    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"codec_test_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
     payload = {
         "task": task,
         "checkpoint": ckpt,
         "config": args.config,
-        "result": result,
+        "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)
+        json.dump(payload, f, indent=2, default=str)
     print(f"Wrote {out_json}")
     return 0